Update app.py

app.py

@@ -52,7 +52,7 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 
 # Use the base model's ID
 base_model_id = "stabilityai/stablelm-3b-4e1t"
-model_directory = "vaishakgkumar/stablemedv3"
+model_directory = "vaishakgkumar/stablemedv1"
 # Instantiate the Tokenizer
 tokenizer = AutoTokenizer.from_pretrained("stabilityai/stablelm-3b-4e1t", token=hf_token, trust_remote_code=True, padding_side="left")
 # tokenizer = AutoTokenizer.from_pretrained("vaishakgkumar/stablemedv3", trust_remote_code=True, padding_side="left")
@@ -60,9 +60,9 @@ tokenizer.pad_token = tokenizer.eos_token
 tokenizer.padding_side = 'left'
 
 # Load the PEFT model
-peft_config = PeftConfig.from_pretrained("vaishakgkumar/stablemedv3", token=hf_token)
+peft_config = PeftConfig.from_pretrained("vaishakgkumar/stablemedv1", token=hf_token)
 peft_model = AutoModelForCausalLM.from_pretrained("stabilityai/stablelm-3b-4e1t", token=hf_token, trust_remote_code=True)
-peft_model = PeftModel.from_pretrained(peft_model, "vaishakgkumar/stablemedv3", token=hf_token)
+peft_model = PeftModel.from_pretrained(peft_model, "vaishakgkumar/stablemedv1", token=hf_token)
 
 class ChatBot:
     def __init__(self):
@@ -82,7 +82,7 @@ class ChatBot:
             chat_history_ids = user_input_ids
 
         # Generate a response using the PEFT model
-        response = peft_model.generate(input_ids=chat_history_ids, max_length=1200, pad_token_id=tokenizer.eos_token_id)
+        response = peft_model.generate(input_ids=chat_history_ids, max_length=512, pad_token_id=tokenizer.eos_token_id)
 
         # Update chat history
         self.history = chat_history_ids