Update app.py

app.py

@@ -34,7 +34,7 @@ def generate_text(user_prompt):
   # Generate text with the complete prompt and limit the maximum length to 256 tokens
   output = model.generate(
       input_ids=prompt_encoded,
-      max_length=256,
+      max_length=1550,
       num_beams=1,
       num_return_sequences=1,  
       do_sample=True,  
@@ -48,9 +48,11 @@ def generate_text(user_prompt):
   generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
 
   # Extract the assistant's response
-  assistant_response = generated_text.split("</s>")[-1]
+  assistant_response = generated_text.split("<|user|>")[-1]
   assistant_response = assistant_response.replace(f"{user_prompt}", "").strip()
   assistant_response = assistant_response.replace(system, "").strip()
+  assistant_response = assistant_response.replace("<|system|>", "").strip()
+  assistant_response = assistant_response.replace("<|assistant|>", "").strip()
   
   return assistant_response