"Update space"

Signed-off-by: Emmanuel Frimpong Asante <[email protected]>

app.py

@@ -105,9 +105,17 @@ generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
 # Define Mistral-based response generation with streaming support
 def mistral_response_stream(user_input):
     try:
-        inputs = tokenizer(user_input, return_tensors="pt", truncation=True)
-        outputs = model.generate(inputs["input_ids"], max_length=150, do_sample=True, temperature=0.7,
-                                 pad_token_id=tokenizer.eos_token_id)
+        # Set truncation and max_length explicitly and pass attention_mask
+        inputs = tokenizer(user_input, return_tensors="pt", truncation=True, max_length=150, padding=True)
+        attention_mask = inputs["attention_mask"]
+        outputs = model.generate(
+            inputs["input_ids"],
+            max_length=150,
+            do_sample=True,
+            temperature=0.7,
+            pad_token_id=tokenizer.eos_token_id,
+            attention_mask=attention_mask
+        )
 
         # Stream the response token by token
         for output in outputs: