Update translators.py

translators.py

@@ -13,7 +13,7 @@ class Translators:
         return response.json()[0][0][0]
     
     def riva4B11(self):
-        set language_pairs = {
+        language_pairs = {
           'en-zh-cn': {'source': 'English', 'target': 'Simplified Chinese'},
           'en-zh': {'source': 'English', 'target': 'Simplified Chinese'},
           'en-zh-tw': {'source': 'English', 'target': 'Traditional Chinese'},
@@ -43,30 +43,29 @@ class Translators:
           'pt-en': {'source': 'Brazilian Portuguese', 'target': 'English'},
           'pt-br-en': {'source': 'Brazilian Portuguese', 'target': 'English'},
         }
-        tokenizer = AutoTokenizer.from_pretrained(self.model)
-        model = AutoModelForCausalLM.from_pretrained(self.model,
+        tokenizer = AutoTokenizer.from_pretrained(self.model_name)
+        model = AutoModelForCausalLM.from_pretrained(self.model_name,
                 torch_dtype="auto", device_map="auto")
         # Use the prompt template (along with chat template)
-        messages = [
-            {
+        messages = [{
                 "role": "system",
                 "content": f"{self.sl}-{self.tl}",
             },
-            {"role": "user", "content": self.input,
-         ]
+            {"role": "user", "content": self.input_text}]
         tokenized_chat = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt").to(model.device)
-        outputs = model.generate(tokenized_chat,  max_new_tokens=128, pad_token_id=tokenizer.eos_token_id)
-        return tokenizer.decode(outputs[0]))
+        outputs = model.generate(tokenized_chat,  max_new_tokens=self.max_new_tokens, pad_token_id=tokenizer.eos_token_id)
+        return tokenizer.decode(outputs[0]).split('<s>Assistant')[1].strip().removesuffix('</s>')
     
     def translategemma(self):
         from huggingface_hub import login
-        login(token=os.environ.get("HF_TOKEN"))
+        hftoken=os.environ.get("HF_TOKEN")
+        login(token=hftoken)
         pipe = pipeline(
             "image-text-to-text",
             model = "google/translategemma-4b-it",
             device = self.device,
-            dtype = torch.bfloat16
-        )      
+            dtype = torch.bfloat16,
+            token=hftoken)      
         # ---- Text Translation ----
         messages = [
             {
@@ -555,4 +554,4 @@ class Translators:
             message_text = f"Translated from {sl} to {tl} with Bergamot {model_name}."
         except Exception as error:
             response = error
-        return translated_text, message_text
+        return translated_text, message_text