updates

app.py

@@ -4,20 +4,16 @@ import torch
 from gradio import State
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
 # Load the tokenizer and model
-tokenizer = AutoTokenizer.from_pretrained("berkeley-nest/Starling-LM-7B-alpha")
-model = AutoModelForCausalLM.from_pretrained("berkeley-nest/Starling-LM-7B-alpha")
+tokenizer = AutoTokenizer.from_pretrained("berkeley-nest/Starling-LM-7B-alpha").to(device)
+model = AutoModelForCausalLM.from_pretrained("berkeley-nest/Starling-LM-7B-alpha").to(device)
 
-# Ensure the model is in evaluation mode
-model.eval()
-
-# Move model to GPU if available
-if torch.cuda.is_available():
-    model = model.to("cuda").half()
 
 @spaces.GPU
 def generate_response(user_input, chat_history):
+    
     prompt = "GPT4 Correct User: " + user_input + "GPT4 Correct Assistant: "
     if chat_history:
         prompt = chat_history + prompt
@@ -28,7 +24,7 @@ def generate_response(user_input, chat_history):
 
     with torch.no_grad():
         # Generate the model's output
-        output = model.generate(**inputs, max_length=1024, num_return_sequences=1, pad_token_id=tokenizer.eos_token_id)
+        output = model.generate(**inputs, max_length=512, num_return_sequences=1, pad_token_id=tokenizer.eos_token_id)
     response = tokenizer.decode(output[0], skip_special_tokens=True)
 
     # Update chat history