Update app.py

app.py

@@ -1,38 +1,63 @@
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+MODEL_NAME = "Devishetty100/savyasachi"
+
+# Load tokenizer
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 
-# Load the fine-tuned Savyasachi model
-model_name = "Devishetty100/savyasachi"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name).to("cuda" if torch.cuda.is_available() else "cpu")
+# Load model (SAFE for Spaces)
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_NAME,
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+    device_map="auto",
+)
 
-# Generate response
-def chat(user_input, max_new_tokens=150, temperature=1.0):
-    input_ids = tokenizer(user_input, return_tensors="pt").to(model.device).input_ids
-    output_ids = model.generate(
-        input_ids,
+# Chat function
+def chat(user_input, history, max_new_tokens=200, temperature=1.0):
+    messages = []
+
+    for user, assistant in history:
+        messages.append({"role": "user", "content": user})
+        messages.append({"role": "assistant", "content": assistant})
+
+    messages.append({"role": "user", "content": user_input})
+
+    prompt = tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=True,
+    )
+
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+
+    outputs = model.generate(
+        **inputs,
         max_new_tokens=max_new_tokens,
         temperature=temperature,
         top_p=0.95,
         top_k=64,
         do_sample=True,
-        eos_token_id=tokenizer.eos_token_id
+        eos_token_id=tokenizer.eos_token_id,
     )
-    return tokenizer.decode(output_ids[0], skip_special_tokens=True)
 
-# Gradio interface
+    response = tokenizer.decode(outputs[0][inputs["input_ids"].shape[-1]:], skip_special_tokens=True)
+    return response
+
+# Gradio UI
 with gr.Blocks() as demo:
-    gr.Markdown("# Chat with Savyasachi – Devotee of Lord Krishna")
+    gr.Markdown("## 🕉️ Savyasachi — Devotee of Lord Krishna")
+
     chatbot = gr.Chatbot()
-    user_input = gr.Textbox(label="You")
-    submit_btn = gr.Button("Send")
-    
-    def respond(user_message, chat_history):
-        response = chat(user_message)
-        chat_history.append((user_message, response))
-        return chat_history, ""
-    
-    submit_btn.click(respond, [user_input, chatbot], [chatbot, user_input])
-    
+    user_input = gr.Textbox(label="Ask Krishna")
+    send = gr.Button("Send")
+
+    def respond(message, history):
+        reply = chat(message, history)
+        history.append((message, reply))
+        return history, ""
+
+    send.click(respond, [user_input, chatbot], [chatbot, user_input])
+
 demo.launch()