Update app.py

app.py

@@ -2,13 +2,22 @@ import streamlit as st
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-# Load model and tokenizer
-model = AutoModelForCausalLM.from_pretrained("adi2606/MenstrualQA")
+# Set up the device to use CPU only
+device = torch.device("cpu")
+
+# Load model and tokenizer, then move the model to the appropriate device
+model = AutoModelForCausalLM.from_pretrained("adi2606/MenstrualQA").to(device)
 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", padding=True, truncation=True)
+def generate_response(message: str, temperature: float = 0.4, repetition_penalty: float = 1.1, max_input_length: int = 256) -> str:
+    inputs = tokenizer(
+        message, 
+        return_tensors="pt", 
+        padding=True, 
+        truncation=True, 
+        max_length=max_input_length
+    ).to(device)
 
     # Generate the response
     output = model.generate(