Update app.py

app.py

@@ -22,7 +22,6 @@ def load_model():
     
     tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
     
-    # Create pipeline
     pipe = pipeline(
         "text-generation",
         model=model,
@@ -36,24 +35,38 @@ pipe = load_model()
 
 @spaces.GPU(duration=110)
 def generate_response(prompt, max_length=1024):
-    # Convert prompt into messages format
+    # Create messages with system prompt
     messages = [
-        {"role": "system", "content": "You are a helpful AI asistent. You always think step by step."},
+        {"role": "system", "content": "You are a helpful AI assistant. You always think step by step."},
         {"role": "user", "content": prompt}
     ]
     
     # Generate response using pipeline
     outputs = pipe(messages, max_new_tokens=max_length)
     
-    # Extract the generated text
+    # Extract the generated text - outputs[0] is a dict with 'generated_text'
     response = outputs[0]["generated_text"]
     
-    # Since pipeline returns the full conversation, we want to extract just the response
-    # Split by the prompt and take the last part
-    response_only = response.split(prompt)[-1].strip()
+    # Find the user's prompt in the response and get everything after it
+    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
     
     return response_only
 
+# Example with proper line breaks
+example_prompt = """For a scalar field theory with interaction Lagrangian $\mathcal{L}_{int} = g\phi^3 + \lambda\phi^4$:
+
+1. Enumerate all possible 1-loop Feynman diagrams contributing to the scalar propagator
+
+2. For each diagram, write down its loop contribution
+
+3. Provide Mathematica code to calculate these loop amplitudes with dimensional regularization at $d=4-\epsilon$
+
+Please explain your reasoning step by step."""
+
 demo = gr.Interface(
     fn=generate_response,
     inputs=[
@@ -69,7 +82,7 @@ demo = gr.Interface(
     
 Model: [benhaotang/phi4-qwq-sky-t1]({MODEL_URL})""",
     examples=[
-        ["For a scalar field theory with interaction Lagrangian $\mathcal{L}_{int} = g\phi^3 + \lambda\phi^4$:\n 1.Enumerate all possible 1-loop Feynman diagrams contributing to a 2-to-2 scattering process\n2.For each diagram, write down its corresponding amplitude\n3. Provide Mathematica code to calculate these loop amplitudes\n Please explain your reasoning step by step."]
+        [example_prompt]  # Now using the formatted example
     ]
 )