Update app.py

app.py

@@ -6,22 +6,20 @@ model_name = "mistralai/Mistral-7B-Instruct-v0.2"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(model_name)
 
-def generate_response(messages):
+def generate_response(input_text):
     """
     Generate response based on the given user messages.
     Parameters:
-    - messages (list): A list of dictionaries containing user messages with roles.
+    - input_text (str): A single string containing all user messages.
     Returns:
     - response (str): The generated response.
     """
-    # Apply chat template and encode messages
-    encodeds = tokenizer.apply_chat_template(messages, return_tensors="pt")
-    # Move inputs to device
-    model_inputs = encodeds 
+    # Tokenize the input text
+    inputs = tokenizer.encode(input_text, return_tensors="pt", max_length=512, truncation=True)
     # Generate response
-    generated_ids = model.generate(model_inputs, max_new_tokens=1000, do_sample=True)
+    generated_ids = model.generate(inputs, max_length=1024, do_sample=True)
     # Decode the generated response
-    response = tokenizer.batch_decode(generated_ids)[0]
+    response = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
     return response
 
 # Define Gradio interface components