Update main.py

main.py

@@ -19,22 +19,26 @@ class PlantChatbot:
     def __init__(self, preprocessed_data_path: str):
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
         
+        # Load the sentence transformer model from local path
         self.embedding_model = SentenceTransformer(
-            model_name_or_path="sentence-transformers/all-mpnet-base-v2",
+            model_name_or_path="/app/models/sentence_transformer",
             device=self.device
         )
         
         self.load_data(preprocessed_data_path)
         
+        # Load the tokenizer and model from local path
         self.tokenizer = AutoTokenizer.from_pretrained(
-            "Qwen/Qwen2.5-1.5B-Instruct",
-            trust_remote_code=True
+            "/app/models/qwen",
+            trust_remote_code=True,
+            local_files_only=True
         )
         self.model = AutoModelForCausalLM.from_pretrained(
-            "Qwen/Qwen2.5-1.5B-Instruct",
-            device_map="auto",
+            "/app/models/qwen",
+            device_map="auto" if torch.cuda.is_available() else None,
             trust_remote_code=True,
-            torch_dtype=torch.float16
+            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+            local_files_only=True
         )
         
         if self.tokenizer.pad_token is None: