handle dict

app.py

@@ -16,10 +16,12 @@ def respond(message, history, system_message, max_tokens, temperature, top_p):
     prompt += f"\nUser: {message}\nAssistant:"
 
     # Generate a response using the model
-    response = text_generator(prompt, max_length=max_tokens, temperature=temperature, top_p=top_p, do_sample=True)
-    
-    # Extract the assistant's response part (after "Assistant:")
-    assistant_response = response.split("Assistant:", 1)[-1].strip()
+    response = text_generator(prompt, max_length=max_tokens, temperature=temperature, top_p=top_p, do_sample=True, truncation=True)
+
+    # Extract the generated text from the response list
+    assistant_response = response[0]['generated_text']
+    # Optionally trim the assistant response if it includes the prompt again
+    assistant_response = assistant_response.split("Assistant:", 1)[-1].strip()
     return assistant_response
 
 """