Create carebridge_client.py

carebridge_client.py

@@ -0,0 +1,70 @@
+import os
+import torch
+from transformers import AutoModelForImageTextToText, AutoProcessor
+from PIL import Image
+import spaces
+import librosa
+from gtts import gTTS
+import tempfile
+
+class CareBridgeTranslator:
+    def __init__(self, model_id="google/translategemma-4b-it", device=None):
+        self.model_id = model_id
+        if device is None:
+            self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        else:
+            self.device = device
+        print(f"[SIMBOTI] Loading model {model_id}...")
+        self.processor = AutoProcessor.from_pretrained(model_id)
+        self.model = AutoModelForImageTextToText.from_pretrained(model_id, device_map=self.device, torch_dtype=torch.float16 if self.device == "cuda" else torch.float32)
+        print("[SIMBOTI] Model loaded successfully.")
+        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"}
+
+    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 f"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)
+        tgt_code = self.LANG_MAP.get(target_lang_name)
+        if not src_code or not tgt_code:
+            return "Error: Language not supported."
+        if isinstance(image_path, str):
+            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)
+        tgt_code = self.LANG_MAP.get(target_lang_name)
+        if not src_code or not tgt_code:
+            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")
+        try:
+            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:
+            print(f"TTS Error: {e}")
+            return None
+
+    @spaces.GPU()
+    def _run_inference(self, messages):
+        inputs = self.processor.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_dict=True, return_tensors="pt").to(self.device)
+        with torch.no_grad():
+            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()