Update app.py

app.py

@@ -41,20 +41,27 @@ def generate_response(prompt, max_length=1024):
         {"role": "user", "content": prompt}
     ]
     
-    # Generate response using pipeline
     outputs = pipe(messages, max_new_tokens=max_length)
     
-    # Extract the generated text - outputs[0] is a dict with 'generated_text'
-    response = outputs[0]["generated_text"]
-    
-    # Find the user's prompt in the response and get everything after it
+    # Find assistant's response in the output
     try:
-        start_idx = response.find(prompt) + len(prompt)
-        response_only = response[start_idx:].strip()
-    except:
-        response_only = response  # Fallback to full response if splitting fails
+        # The output contains the full conversation history
+        generated_text = outputs[0]["generated_text"]
+        # Look for the last assistant message
+        assistant_prefix = "{'role': 'assistant', 'content': '"
+        assistant_start = generated_text.rfind(assistant_prefix)
+        if assistant_start != -1:
+            # Move past the prefix
+            content_start = assistant_start + len(assistant_prefix)
+            # Find the end of the content (before the closing quote and brace)
+            content_end = generated_text.rfind("'}")
+            if content_end != -1:
+                return generated_text[content_start:content_end]
+    except Exception as e:
+        print(f"Error extracting response: {e}")
     
-    return response_only
+    # Fallback: return the raw generated text if extraction fails
+    return outputs[0]["generated_text"]
 
 # Example with proper line breaks
 example_prompt = """For a scalar field theory with interaction Lagrangian $\mathcal{L}_{int} = g\phi^3 + \lambda\phi^4$:
@@ -78,7 +85,7 @@ demo = gr.Interface(
     ],
     outputs=gr.Textbox(label="Response", lines=10),
     title="benhaotang/phi4-qwq-sky-t1",
-    description=f""" To achieve CoT and science reasoning on small scale
+    description=f""" To achieve CoT and science reasoning on small scale with a merge of CoT finetuned phi4 model.
     
 Model: [benhaotang/phi4-qwq-sky-t1]({MODEL_URL})""",
     examples=[