Update app.py

app.py

@@ -8,6 +8,8 @@ model_name = "CreitinGameplays/bloom-3b-conversational"
 # Load tokenizer and model
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(model_name)
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)
 
 def generate_text(user_prompt):
   """Generates text using the BLOOM model from Hugging Face Transformers and removes the user prompt."""
@@ -15,7 +17,7 @@ def generate_text(user_prompt):
   prompt = f"<|system|> You are a helpful AI assistant. </s> <|prompter|> {user_prompt} </s> <|assistant|>"
 
   # Encode the entire prompt into tokens
-  prompt_encoded = tokenizer(prompt, return_tensors="pt").input_ids
+  prompt_encoded = tokenizer.encode(prompt, return_tensors="pt").to(device)
 
   # Generate text with the complete prompt and limit the maximum length to 256 tokens
   output = model.generate(