Update app.py

app.py

@@ -48,10 +48,18 @@ def load_model(model_name):
             del model
             torch.cuda.empty_cache()
         
+        # Check if the model is in models or quantized_models and load accordingly
+        if model_name in models:
+            model_path = models[model_name]
+        elif model_name in quantized_models:
+            model_path = quantized_models[model_name]
+        else:
+            raise ValueError(f"Model {model_name} not found in either models or quantized_models.")
+
         tokenizer = AutoTokenizer.from_pretrained(
-            models[model_name], 
+            model_path, 
             padding_side='left', 
-            token=hugging_face_token, 
+            use_auth_token=hugging_face_token, 
             trust_remote_code=True
         )
         
@@ -61,8 +69,8 @@ def load_model(model_name):
             tokenizer.add_special_tokens({'pad_token': tokenizer.eos_token})
         
         model = AutoModelForCausalLM.from_pretrained(
-            models[model_name], 
-            token=hugging_face_token, 
+            model_path, 
+            use_auth_token=hugging_face_token, 
             trust_remote_code=True
         )
         
@@ -76,6 +84,7 @@ def load_model(model_name):
         raise e
     return tokenizer, model
 
+
 # Ensure the initial model is loaded
 tokenizer, model = load_model(selected_model)