Update app.py

app.py

@@ -78,10 +78,11 @@ def generate_response(system_instruction, user_input):
     # Generate model response
     with torch.no_grad():
         generated_ids = model.generate(model_inputs, max_new_tokens=1000, do_sample=True)
-    # Extract only the bot's response, omitting the prompt
-    response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0] #.split("\n")[-2:]
+    # Find everything after the <|assistant|> tag
+    decoded_output = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
+    assistant_response = decoded_output.split("<|assistant|>")[-1].strip()
 
-    return response
+    return assistant_response
 
 # Gradio interface setup
 with gr.Blocks() as demo: