Update modules/translation_model.py

modules/translation_model.py

@@ -4,27 +4,34 @@ import logging
 
 class TranslationModel:
     def __init__(self, cache_dir="models/"):
-        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-        if self.device.type == "cpu":
-            logging.warning("GPU not found, using CPU. Translation will be slower.")
+        self.device = torch.device("cpu")
+        logging.info("Using CPU for translations")
         
         self.model_name = "facebook/m2m100_1.2B"
-        self.tokenizer = M2M100Tokenizer.from_pretrained(self.model_name, cache_dir=cache_dir)
+        self.tokenizer = M2M100Tokenizer.from_pretrained(
+            self.model_name, 
+            cache_dir=cache_dir,
+            local_files_only=True  # Only use cached files
+        )
         self.model = M2M100ForConditionalGeneration.from_pretrained(
             self.model_name, 
-            cache_dir=cache_dir
-        ).to(self.device)
+            cache_dir=cache_dir,
+            local_files_only=True,
+            device_map="cpu",
+            low_cpu_mem_usage=True
+        )
         self.model.eval()
 
     def translate(self, text: str, source_lang: str, target_lang: str) -> str:
         try:
             self.tokenizer.src_lang = source_lang
-            encoded = self.tokenizer(text, return_tensors="pt").to(self.device)
+            encoded = self.tokenizer(text, return_tensors="pt")
             
             with torch.no_grad():
                 generated = self.model.generate(
                     **encoded,
-                    forced_bos_token_id=self.tokenizer.get_lang_id(target_lang)
+                    forced_bos_token_id=self.tokenizer.get_lang_id(target_lang),
+                    max_length=128
                 )
                 
             return self.tokenizer.batch_decode(generated, skip_special_tokens=True)[0]