Update main.py

main.py

@@ -19,26 +19,35 @@ class PlantChatbot:
     def __init__(self, preprocessed_data_path: str):
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
         
+        print("Loading sentence transformer model...")
         self.embedding_model = SentenceTransformer(
-            model_name_or_path="all-mpnet-base-v2",
+            model_name_or_path="/app/models/sentence_transformer",
             device=self.device
         )
         
+        print("Loading data...")
         self.load_data(preprocessed_data_path)
         
+        print("Loading Qwen tokenizer...")
         self.tokenizer = AutoTokenizer.from_pretrained(
-            "Qwen/Qwen2.5-0.5B-Instruct",
-            trust_remote_code=True
+            "/app/models/qwen",
+            trust_remote_code=True,
+            local_files_only=True
         )
+        
+        print("Loading Qwen model...")
         self.model = AutoModelForCausalLM.from_pretrained(
-            "Qwen/Qwen2.5-0.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:
             self.tokenizer.pad_token = self.tokenizer.eos_token
+        
+        print("Initialization complete!")
 
     def load_data(self, preprocessed_data_path: str):
         df = pd.read_csv(preprocessed_data_path)