more

app.py

@@ -33,20 +33,23 @@ def chat_with_model(user_message, chat_history=None):
         ).to("cuda")
 
         # Generate response
-        text_streamer = TextStreamer(tokenizer, skip_prompt=True)
-        output = model.generate(
-            input_ids=inputs,
-            streamer=text_streamer,
+        output_ids = model.generate(
+            input_ids=inputs["input_ids"],
+            attention_mask=inputs["attention_mask"],  # Ensure attention_mask is included
+            streamer=None,  # Collect output as tensor
             max_new_tokens=128,
             use_cache=True,
             temperature=1.5,
             min_p=0.1,
         )
 
+        # Decode the generated tokens into a string
+        response = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+
         # Append the response to the chat history
         if chat_history is None:
             chat_history = []
-        chat_history.append((user_message, output))
+        chat_history.append((user_message, response))
         return "", chat_history
     except Exception as e:
         return f"Error: {str(e)}", chat_history