Update app.py

app.py

@@ -60,10 +60,33 @@ def progress_bar_html(label: str) -> str:
 </style>
     '''
 
+# Helper function for the chat interface
+def apply_chat_template_for_text(conversation, add_generation_prompt=True):
+    """
+    Concatenates a conversation (list of dict with keys "role" and "content") 
+    into a single string prompt. If add_generation_prompt is True, appends "assistant:".
+    """
+    prompt = ""
+    for msg in conversation:
+        prompt += f"{msg['role']}: {msg['content']}\n"
+    if add_generation_prompt:
+        prompt += "assistant:"
+    return prompt
+
+def clean_chat_history(chat_history):
+    """
+    Filter out any chat entries whose "content" is not a string.
+    """
+    cleaned = []
+    for msg in chat_history:
+        if isinstance(msg, dict) and isinstance(msg.get("content"), str):
+            cleaned.append(msg)
+    return cleaned
+
 # ---------------------------
 # 1. Chat Interface Tab
 # ---------------------------
-# Uses a text-only model: FastThink-0.5B-Tiny
+# Uses a text-only model: DeepHermes-3-Llama-3-3B-Preview-abliterated
 
 model_id_text = "prithivMLmods/DeepHermes-3-Llama-3-3B-Preview-abliterated"
 tokenizer = AutoTokenizer.from_pretrained(model_id_text)
@@ -74,16 +97,6 @@ model = AutoModelForCausalLM.from_pretrained(
 )
 model.eval()
 
-def clean_chat_history(chat_history):
-    """
-    Filter out any chat entries whose "content" is not a string.
-    """
-    cleaned = []
-    for msg in chat_history:
-        if isinstance(msg, dict) and isinstance(msg.get("content"), str):
-            cleaned.append(msg)
-    return cleaned
-
 @spaces.GPU
 def chat_generate(input_text: str, chat_history: list, max_new_tokens: int, temperature: float, top_p: float, top_k: int, repetition_penalty: float):
     """
@@ -92,7 +105,10 @@ def chat_generate(input_text: str, chat_history: list, max_new_tokens: int, temp
     # Prepare conversation by cleaning history and appending the new user message.
     conversation = clean_chat_history(chat_history)
     conversation.append({"role": "user", "content": input_text})
-    input_ids = tokenizer.apply_chat_template(conversation, add_generation_prompt=True, return_tensors="pt")
+    
+    # Instead of tokenizer.apply_chat_template, we use our helper to generate a prompt.
+    prompt_text = apply_chat_template_for_text(conversation, add_generation_prompt=True)
+    input_ids = tokenizer(prompt_text, return_tensors="pt").input_ids
     if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH:
         input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]
     input_ids = input_ids.to(model.device)
@@ -115,7 +131,7 @@ def chat_generate(input_text: str, chat_history: list, max_new_tokens: int, temp
     for new_text in streamer:
         outputs.append(new_text)
     final_response = "".join(outputs)
-    # Append assistant reply to chat history.
+    # Append assistant reply to conversation.
     updated_history = conversation + [{"role": "assistant", "content": final_response}]
     return final_response, updated_history
 
@@ -144,7 +160,7 @@ def generate_qwen_ocr(input_text: str, image):
         {"role": "system", "content": [{"type": "text", "text": "You are a helpful assistant."}]},
         {"role": "user", "content": [{"type": "text", "text": input_text}, {"type": "image", "image": image}]}
     ]
-    # Apply chat template.
+    # Use the processor's chat template.
     prompt_full = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
     inputs = processor(text=[prompt_full], images=[image], return_tensors="pt", padding=True).to("cuda")
     streamer = TextIteratorStreamer(processor, skip_prompt=True, skip_special_tokens=True)