Update app.py

app.py

@@ -25,9 +25,7 @@ model = PeftModel.from_pretrained(model, model_id)
 def greet(text):
     with torch.no_grad():
         # Include EOS token for better context
-        input_text = "<s>### Instruction:\nYou are a data analyst tasked with helping students finding resources, respond in JSON format.\n\n" +
-                     f"### Input:\n{text}\n\n" +
-                     "### Response:\n"
+        input_text = "<s>### User:\n{text}\n\n### Assistant:\n"
         
         batch = tokenizer(input_text, return_tensors='pt', add_special_tokens=True).to(device)
 
@@ -42,8 +40,8 @@ def greet(text):
         # Decode only the generated tokens
         response = tokenizer.decode(output_tokens[0][len(batch['input_ids'][0]):], skip_special_tokens=True)
 
-        # Additional stopping condition at next "\n### Response:"
-        response_parts = response.split("\n### Response:")
+        # Additional stopping condition at next "### Response:"
+        response_parts = response.split("### Assistant:")
         return response_parts[0]  # Return only the first part
 
 iface = gr.Interface(fn=greet, inputs="text", outputs="text", title="PEFT Model for Big Brain")