Update off_topic.py

off_topic.py

@@ -10,28 +10,71 @@ import numpy as np
 import torch
 import PIL
 import imagehash
-from transformers import CLIPModel, CLIPProcessor
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, CLIPModel, CLIPProcessor
 from PIL import Image
 
 
+class Translator:
+    def __init__(self, model_id: str, device: Optional[str] = None):
+        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+        self.model_id = model_id
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            model_id)
+        self.model = AutoModelForSeq2SeqLM.from_pretrained(model_id).to(self.device)
+        self.bos_token_map = self.tokenizer.get_lang_id if hasattr(self.tokenizer, "get_lang_id") else self.tokenizer.lang_code_to_id
+
+    @property
+    def _language_code_mapper(self):
+        if "nllb" in self.model_id.lower():
+            return {"en": "eng_Latn",
+                    "es": "spa_Latn",
+                    "pt": "por_Latn"}
+        elif "m2m" in self.model_id.lower():
+            return {"en": "en",
+                    "es": "es",
+                    "pt": "pt"}
+
+    def translate(self, texts: List[str], src_lang: str, dest_lang: str = "en", max_length: int = 100):
+        self.tokenizer.src_lang = self._language_code_mapper[src_lang]
+        inputs = self.tokenizer(texts, return_tensors="pt").to(self.device)
+        translated_tokens = self.model.generate(
+            **inputs, forced_bos_token_id=self.bos_token_map["eng_Latn"], max_length=max_length
+        )
+        return self.tokenizer.batch_decode(translated_tokens, skip_special_tokens=True)
+
+
 class OffTopicDetector:
-    def __init__(self, model_id: str, device: Optional[str] = None, image_size: str = "E"):
+    def __init__(self, model_id: str, device: Optional[str] = None, image_size: str = "E", translator: Optional[Translator] = None):
         self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
         self.processor = CLIPProcessor.from_pretrained(model_id)
         self.model = CLIPModel.from_pretrained(model_id).to(self.device)
         self.image_size = image_size
+        self.translator = translator
 
     def predict_probas(self, images: List[PIL.Image.Image], domain: str,
+                       title: Optional[str] = None,
                 valid_templates: Optional[List[str]] = None,
                 invalid_classes: Optional[List[str]] = None,
                 autocast: bool = True):
+        site, domain = domain.split("-")
+        domain = re.sub("_", " ", domain).lower()
         if valid_templates:
             valid_classes = [template.format(domain) for template in valid_templates]
         else:
-            valid_classes = [f"a photo of {domain}", f"brochure with {domain} image", f"instructions for {domain}", f"{domain} diagram",
-                             f"{domain} packaging", f"box of {domain}"]
+            valid_classes = [f"a photo of {domain}", f"brochure with {domain} image", f"instructions for {domain}", f"{domain} diagram"]
+        if title:
+            if site == "CBT":
+                translated_title = title
+            else:
+                if site == "MLB":
+                    src_lang = "pt"
+                else:
+                    src_lang = "es"
+                translated_title = self.translator.translate(title, src_lang=src_lang, dest_lang="en", max_length=100)[0]
+            valid_classes.append(translated_title)
         if not invalid_classes:
             invalid_classes = ["promotional ad with store information", "promotional text", "google maps screenshot", "business card", "qr code"]
+        
         n_valid = len(valid_classes)
         classes = valid_classes + invalid_classes
         print(f"Valid classes: {valid_classes}", f"Invalid classes: {invalid_classes}", sep="\n")
@@ -59,18 +102,21 @@ class OffTopicDetector:
         return probas, valid_probas, invalid_probas
 
     def predict_probas_url(self, img_urls: List[str], domain: str,
+                           title: Optional[str] = None,
                 valid_templates: Optional[List[str]] = None,
                 invalid_classes: Optional[List[str]] = None,
                 autocast: bool = True):
         images = self.get_images(img_urls)
         dedup_images = self._filter_dups(images)
-        return self.predict_probas(images, domain, valid_templates, invalid_classes, autocast)
+        return self.predict_probas(images, domain, title, valid_templates, invalid_classes, autocast)
 
     def predict_probas_item(self, url_or_id: str,
+                            use_title: bool = False,
                 valid_templates: Optional[List[str]] = None,
                 invalid_classes: Optional[List[str]] = None):
-        images, domain = self.get_item_data(url_or_id)
-        probas, valid_probas, invalid_probas = self.predict_probas(images, domain, valid_templates,
+        images, domain, title = self.get_item_data(url_or_id)
+        title = title if use_title else None
+        probas, valid_probas, invalid_probas = self.predict_probas(images, domain, title, valid_templates,
                                                             invalid_classes)
         return images, domain, probas, valid_probas, invalid_probas
 
@@ -84,7 +130,8 @@ class OffTopicDetector:
             item_id = re.sub("-", "", url_or_id)
         start = time.time()
         response = httpx.get(f"https://api.mercadolibre.com/items/{item_id}").json()
-        domain = re.sub("_", " ", response["domain_id"].split("-")[-1]).lower()
+        domain = response["domain_id"]
+        title = response["title"]
         img_urls = [x["url"] for x in response["pictures"]]
         img_urls = [x.replace("-O.jpg", f"-{self.image_size}.jpg") for x in img_urls]
         end = time.time()
@@ -92,7 +139,7 @@ class OffTopicDetector:
         print(f"Items API time: {round(duration * 1000, 0)} ms")
         images = self.get_images(img_urls)
         dedup_images = self._filter_dups(images)
-        return dedup_images, domain
+        return dedup_images, domain, title
 
     def _filter_dups(self, images: List):
         if len(images) > 1:
@@ -166,4 +213,4 @@ class OffTopicDetector:
         if title:
             fig.suptitle(title)
         fig.tight_layout()
-        return
+        return