Update app.py

app.py

@@ -10,7 +10,7 @@ import requests
 import cloudinary
 import cloudinary.uploader
 import cloudinary.api
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig
 import speech_recognition as sr
 from pydub import AudioSegment
 from happytransformer import HappyTextToText, TTSettings
@@ -187,13 +187,24 @@ retriever = CustomRagRetriever(index, model_embedding)
 # Load tokenizer and LLM model
 def load_tokenizer_and_model():
     print("Loading tokenizer...")
-    tokenizer = AutoTokenizer.from_pretrained(model_dir)
-    print("Tokenizer loaded successfully.")
-
-    print("Loading LLM model...")
-    model_llm = AutoModelForCausalLM.from_pretrained(model_dir).to(device)
-    print("LLM model loaded successfully.")
-    return tokenizer, model_llm
+    try:
+        tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)
+        print("Tokenizer loaded successfully.")
+
+        print("Loading LLM model...")
+        model_config = AutoConfig.from_pretrained(model_dir, trust_remote_code=True)
+        model_llm = AutoModelForCausalLM.from_pretrained(
+            model_dir,
+            config=model_config,
+            trust_remote_code=True,
+            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+            device_map="auto"
+        ).to(device)
+        print("LLM model loaded successfully.")
+        return tokenizer, model_llm
+    except Exception as e:
+        print(f"Error loading model: {str(e)}")
+        raise
 
 tokenizer, model_llm = load_tokenizer_and_model()