replace Helsinki

app.py

@@ -294,18 +294,80 @@ class EventDetector:
         logger.info(f"Initializing models on device: {device}")
         
         """Initialize all models with GPU support"""
-        # Initialize translation model
-        self.translator = pipeline(
-            "translation",
-            model="Helsinki-NLP/opus-mt-ru-en",
-            device=device
-        )
+        from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
         
-        self.rutranslator = pipeline(
-            "translation",
-            model="Helsinki-NLP/opus-mt-en-ru",
-            device=device
+        logger.info("replacing Helsinki-NLP due to conflict with PyTorch version)")
+        self.translator_tokenizer = AutoTokenizer.from_pretrained("facebook/nllb-200-distilled-600M")
+        self.translator_model = AutoModelForSeq2SeqLM.from_pretrained(
+            "facebook/nllb-200-distilled-600M",
+            use_safetensors=True,
+            device_map=device
         )
+        
+        # Create custom translation function for ru→en
+        def translate_ru_en(text_list):
+            if not isinstance(text_list, list):
+                text_list = [text_list]
+            
+            results = []
+            for text in text_list:
+                if not text:
+                    results.append({"translation_text": ""})
+                    continue
+                    
+                # Prepare input
+                inputs = self.translator_tokenizer(text, return_tensors="pt").to(device)
+                inputs["forced_bos_token_id"] = self.translator_tokenizer.lang_code_to_id["eng_Latn"]
+                
+                # Generate translation
+                with torch.no_grad():
+                    outputs = self.translator_model.generate(
+                        **inputs,
+                        forced_bos_token_id=self.translator_tokenizer.lang_code_to_id["eng_Latn"],
+                        max_length=512,
+                    )
+                
+                # Decode and format like Helsinki-NLP output
+                translation = self.translator_tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]
+                results.append({"translation_text": translation})
+            
+            return results
+        
+        # Create custom translation function for en→ru
+        def translate_en_ru(text_list):
+            if not isinstance(text_list, list):
+                text_list = [text_list]
+            
+            results = []
+            for text in text_list:
+                if not text:
+                    results.append({"translation_text": ""})
+                    continue
+                    
+                # Prepare input
+                inputs = self.translator_tokenizer(text, return_tensors="pt").to(device)
+                inputs["forced_bos_token_id"] = self.translator_tokenizer.lang_code_to_id["rus_Cyrl"]
+                
+                # Generate translation
+                with torch.no_grad():
+                    outputs = self.translator_model.generate(
+                        **inputs,
+                        forced_bos_token_id=self.translator_tokenizer.lang_code_to_id["rus_Cyrl"],
+                        max_length=512,
+                    )
+                
+                # Decode and format like Helsinki-NLP output
+                translation = self.translator_tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]
+                results.append({"translation_text": translation})
+            
+            return results
+        
+        # Replace pipeline with custom functions that mimic the original API
+        self.translator = translate_ru_en
+        self.rutranslator = translate_en_ru
+        
+        logger.info("Translation models replaced successfully!")
+
 
         # Initialize sentiment models
         self.finbert = pipeline(