Fix: Mismatch between qa fn and call

app.py

@@ -51,7 +51,7 @@ def transcribe_audio(audio_filepath):
             prompts=[[{"role": "user", "content": f"Transcribe the following: {model.audio_locator_tag}"}]],
             audios=audio.to(device),
             audio_lens=torch.tensor([audio_lens]).to(device),
-            max_new_tokens=256,
+            max_new_tokens=2048,
         )
 
     # Convert output IDs to text
@@ -60,7 +60,7 @@ def transcribe_audio(audio_filepath):
     return transcript, transcript, initial_message
 
 
-# Simple Q&A function - adapted from working version
+# Simple Q&A function
 @spaces.GPU
 def transcript_qa(transcript, question, history):
     if not transcript:
@@ -69,33 +69,27 @@ def transcript_qa(transcript, question, history):
     if not question:
         return history, ""
     
-    # Add user message to history first
+    # Add user message to history
     history = history + [{"role": "user", "content": question}]
     
     with torch.inference_mode(), model.llm.disable_adapter():
         output_ids = model.generate(
             prompts=[[{"role": "user", "content": f"{question}\n\n{transcript}"}]],
-            max_new_tokens=256,
+            max_new_tokens=2048,
         )
     
-    # Convert output IDs to text and extract answer
-    answer = model.tokenizer.ids_to_text(output_ids[0].cpu())
-    answer = answer.split("<|im_start|>assistant")[-1]
+    ans = model.tokenizer.ids_to_text(output_ids[0].cpu())
+    ans = ans.split("<|im_start|>assistant")[-1]  # get rid of the prompt
     
-    # Remove thinking tags if present
-    if "<think>" in answer:
-        if "</think>" in answer:
-            parts = answer.split("</think>")
-            if len(parts) > 1:
-                answer = parts[-1]  # Get text after thinking
-        else:
-            # If no closing tag, try to get text after opening tag
-            answer = answer.split("<think>")[0]  # Get text before thinking
+    if "<think>" in ans:
+        if "</think>" in ans:
+            ans = ans.split("<think>")[-1]
+            _, ans = ans.split("</think>")  # get rid of the thinking
     
-    answer = answer.strip()
+    ans = ans.strip()
     
     # Add assistant response to history
-    history = history + [{"role": "assistant", "content": answer}]
+    history = history + [{"role": "assistant", "content": ans}]
     
     return history, ""  # Return updated history and clear input
 
@@ -140,6 +134,9 @@ with gr.Blocks(theme=theme) as demo:
         bubble_full_width=False
     )
     
+    def user(user_message, history: list):
+        return "", history + [{"role": "user", "content": user_message}]
+    
     with gr.Row():
         question_input = gr.Textbox(
             label="",