Update app.py

app.py

@@ -12,6 +12,10 @@ model = AutoModelForCausalLM.from_pretrained(
     trust_remote_code=True,
 )
 
+# Ensure pad token exists (many chat models reuse EOS as PAD)
+if tokenizer.pad_token_id is None and tokenizer.eos_token_id is not None:
+    tokenizer.pad_token = tokenizer.eos_token
+
 def _msgs_from_history(history, system_text):
     msgs = []
     if system_text:
@@ -24,23 +28,29 @@ def _msgs_from_history(history, system_text):
     return msgs
 
 def _eos_ids(tok):
-    ids = {tok.eos_token_id}
-    im_end = tok.convert_tokens_to_ids("<|im_end|>")
-    if im_end is not None:
-        ids.add(im_end)
+    # Support ints/lists and optional <|im_end|>
+    ids = set()
+    if tok.eos_token_id is not None:
+        if isinstance(tok.eos_token_id, (list, tuple)):
+            ids.update(tok.eos_token_id)
+        else:
+            ids.add(tok.eos_token_id)
+    try:
+        im_end = tok.convert_tokens_to_ids("<|im_end|>")
+        if im_end is not None and im_end != tok.unk_token_id:
+            ids.add(im_end)
+    except Exception:
+        pass
+    # Fallback: if still empty, just skip setting eos_token_id in GenerationConfig
     return list(ids)
 
-@spaces.GPU()  # REQUIRED for ZeroGPU; remove if using standard GPU hardware
-def gradio_fn(message, history):
-    response = infer_text(history + [(message, None)])
-    return response
-
 def chat_fn(message, history, system_text, temperature, top_p, max_new, min_new):
     msgs = _msgs_from_history(history, system_text) + [{"role": "user", "content": message}]
     prompt = tokenizer.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)
     inputs = tokenizer([prompt], return_tensors="pt").to(model.device)
 
-    gen_cfg = GenerationConfig(
+    eos = _eos_ids(tokenizer)
+    gen_cfg_kwargs = dict(
         do_sample=True,
         temperature=float(temperature),
         top_p=float(top_p),
@@ -48,9 +58,13 @@ def chat_fn(message, history, system_text, temperature, top_p, max_new, min_new)
         min_new_tokens=int(min_new),
         repetition_penalty=1.02,
         no_repeat_ngram_size=3,
-        eos_token_id=_eos_ids(tokenizer),
-        pad_token_id=tokenizer.eos_token_id,
+        pad_token_id=tokenizer.pad_token_id if tokenizer.pad_token_id is not None else tokenizer.eos_token_id,
     )
+    if eos:
+        gen_cfg_kwargs["eos_token_id"] = eos
+
+    gen_cfg = GenerationConfig(**gen_cfg_kwargs)
+
     with torch.no_grad():
         out = model.generate(**inputs, generation_config=gen_cfg)
 
@@ -59,11 +73,17 @@ def chat_fn(message, history, system_text, temperature, top_p, max_new, min_new)
     reply = tokenizer.batch_decode(new_tokens, skip_special_tokens=True)[0].strip()
     return reply
 
+# Wrap for ChatInterface + ZeroGPU
+@spaces.GPU()  # REQUIRED for ZeroGPU; remove if using standard GPU hardware
+def gradio_fn(message, history, system_text, temperature, top_p, max_new, min_new):
+    return chat_fn(message, history, system_text, temperature, top_p, max_new, min_new)
+
 with gr.Blocks() as demo:
     gr.Markdown(
         "<h1 style='text-align:center'>Gita Assistant (Qwen2.5-3B Fine-tuned)</h1>"
         "<p style='text-align:center'>Ask in English / हिंदी / ગુજરાતી. The assistant cites verses when relevant.</p>"
     )
+
     system_box = gr.Textbox(
         value="Reply in the user’s language with 2–3 concise points (200–400 words); cite Gita verses when relevant.",
         label="System prompt",
@@ -75,6 +95,7 @@ with gr.Blocks() as demo:
 
     chat = gr.ChatInterface(
         fn=gradio_fn,
+        additional_inputs=[system_box, temperature, top_p, max_new, min_new],
         examples=[
             "Hello!",
             "How can I overcome fear of failure?",
@@ -82,8 +103,7 @@ with gr.Blocks() as demo:
             "What can I do to stop overthinking?"
         ],
         chatbot=gr.Chatbot(elem_classes="chatbot"),
-        theme="compact",
     )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()