Update app.py

app.py

@@ -2,24 +2,23 @@ import streamlit as st
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-# Set up the device
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
 # Load model and tokenizer
-model = AutoModelForCausalLM.from_pretrained("adi2606/MenstrualQA").to(device)
+model = AutoModelForCausalLM.from_pretrained("adi2606/MenstrualQA")
 tokenizer = AutoTokenizer.from_pretrained("adi2606/MenstrualQA")
 
 # Function to generate a response from the chatbot
 def generate_response(message: str, temperature: float = 0.4, repetition_penalty: float = 1.1) -> str:
-    inputs = tokenizer(message, return_tensors="pt").to(device)
+    inputs = tokenizer(message, return_tensors="pt", padding=True, truncation=True).to(device)
 
     # Generate the response
     output = model.generate(
         inputs['input_ids'],
+        attention_mask=inputs['attention_mask'],
         max_length=512,
         temperature=temperature,
         repetition_penalty=repetition_penalty,
-        do_sample=True
+        do_sample=True,
+        pad_token_id=tokenizer.eos_token_id
     ) 
 
     # Decode the generated output