Update carebridge_client.py

carebridge_client.py

@@ -8,66 +8,66 @@ from gtts import gTTS
 import tempfile
 
 class CareBridgeTranslator:
-        def __init__(self, model_id="google/translategemma-4b-it"):
-                    self.model_id = model_id
-                    self.model = None
-                    self.processor = None
-                    self.LANG_MAP = {"English": "en", "Polish": "pl", "Romanian": "ro", "Punjabi": "pa", "Urdu": "ur", "Portuguese": "pt", "Spanish": "es", "Arabic": "ar", "Bengali": "bn", "Gujarati": "gu", "Italian": "it"}
-                    print("[SIMBOTI] Translator initialized. Model will load on first use.")
+    def __init__(self, model_id="google/translategemma-4b-it"):
+        self.model_id = model_id
+        self.model = None
+        self.processor = None
+        self.LANG_MAP = {"English": "en", "Polish": "pl", "Romanian": "ro", "Punjabi": "pa", "Urdu": "ur", "Portuguese": "pt", "Spanish": "es", "Arabic": "ar", "Bengali": "bn", "Gujarati": "gu", "Italian": "it"}
+        print("[SIMBOTI] Translator initialized. Model will load on first use.")
 
-        def _load_model(self):
-                    if self.model is None:
-                                    print(f"[SIMBOTI] Loading model {self.model_id}...")
-                                    self.processor = AutoProcessor.from_pretrained(self.model_id)
-                                    self.model = AutoModelForImageTextToText.from_pretrained(self.model_id, device_map="cuda", torch_dtype=torch.float16)
-                                    print("[SIMBOTI] Model loaded successfully.")
+    def _load_model(self):
+        if self.model is None:
+            print(f"[SIMBOTI] Loading model ${self.model_id}...")
+            self.processor = AutoProcessor.from_pretrained(self.model_id)
+            self.model = AutoModelForImageTextToText.from_pretrained(self.model_id, device_map="cuda", torch_dtype=torch.float16)
+            print("[SIMBOTI] Model loaded successfully.")
 
-                def translate_text(self, text, source_lang_name, target_lang_name):
-                            src_code = self.LANG_MAP.get(source_lang_name)
-                            tgt_code = self.LANG_MAP.get(target_lang_name)
-                            if not src_code or not tgt_code:
-                                            return "Error: Language not supported."
-                                        message = {"role": "user", "content": [{"type": "text", "source_lang_code": src_code, "target_lang_code": tgt_code, "text": text}]}
+    def translate_text(self, text, source_lang_name, target_lang_name):
+        src_code = self.LANG_MAP.get(source_lang_name)
+        tgt_code = self.LANG_MAP.get(target_lang_name)
+        if not src_code or not tgt_code:
+            return "Error: Language not supported."
+        message = {"role": "user", "content": [{"type": "text", "source_lang_code": src_code, "target_lang_code": tgt_code, "text": text}]}
         return self._run_inference([message])
 
     def translate_image(self, image_path, source_lang_name, target_lang_name):
-                src_code = self.LANG_MAP.get(source_lang_name)
+        src_code = self.LANG_MAP.get(source_lang_name)
         tgt_code = self.LANG_MAP.get(target_lang_name)
         if not src_code or not tgt_code:
-                        return "Error: Language not supported."
+            return "Error: Language not supported."
         if isinstance(image_path, str):
-                        image = Image.open(image_path)
-else:
+            image = Image.open(image_path)
+        else:
             image = image_path
         message = {"role": "user", "content": [{"type": "image", "source_lang_code": src_code, "target_lang_code": tgt_code, "image": image}]}
         return self._run_inference([message])
 
     def translate_audio(self, audio_path, source_lang_name, target_lang_name):
-                src_code = self.LANG_MAP.get(source_lang_name)
+        src_code = self.LANG_MAP.get(source_lang_name)
         tgt_code = self.LANG_MAP.get(target_lang_name)
         if not src_code or not tgt_code:
-                        return "Error: Language not supported."
+            return "Error: Language not supported."
         audio, sr = librosa.load(audio_path, sr=16000)
         message = {"role": "user", "content": [{"type": "audio", "source_lang_code": src_code, "target_lang_code": tgt_code, "audio": audio}]}
         return self._run_inference([message])
 
     def speak_text(self, text, lang_name):
-                lang_code = self.LANG_MAP.get(lang_name, "en")
+        lang_code = self.LANG_MAP.get(lang_name, "en")
         try:
-                        tts = gTTS(text=text, lang=lang_code)
+            tts = gTTS(text=text, lang=lang_code)
             temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".mp3")
             tts.save(temp_file.name)
             return temp_file.name
-except Exception as e:
+        except Exception as e:
             print(f"TTS Error: {e}")
             return None
 
     @spaces.GPU()
     def _run_inference(self, messages):
-                self._load_model()
+        self._load_model()
         inputs = self.processor.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_dict=True, return_tensors="pt").to("cuda")
         with torch.no_grad():
-                        outputs = self.model.generate(**inputs, max_new_tokens=512, do_sample=False)
+            outputs = self.model.generate(**inputs, max_new_tokens=512, do_sample=False)
         input_len = inputs["input_ids"].shape[-1]
         decoded = self.processor.decode(outputs[0][input_len:], skip_special_tokens=True)
-        return decoded.strip()
+        return decoded.strip()