Update app.py

app.py

@@ -1,10 +1,10 @@
 from fastapi import FastAPI, HTTPException
-from transformers import M2M100ForConditionalGeneration, M2M100Tokenizer
+from transformers import pipeline
 import langdetect
 import logging
 import os
 
-# Atur direktori cache untuk Hugging Face
+# Atur direktori cache
 os.environ["HF_HOME"] = "/app/cache"
 os.environ["TRANSFORMERS_CACHE"] = "/app/cache"
 
@@ -13,53 +13,36 @@ app = FastAPI()
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
+# Daftar model untuk bahasa yang didukung
+MODEL_MAP = {
+    "id": "Helsinki-NLP/opus-mt-id-en",  # Indonesia ke Inggris
+    "th": "Helsinki-NLP/opus-mt-th-en",  # Thailand ke Inggris
+    "fr": "Helsinki-NLP/opus-mt-fr-en",  # Prancis ke Inggris
+    "es": "Helsinki-NLP/opus-mt-es-en",  # Spanyol ke Inggris
+    "ja": "Helsinki-NLP/opus-mt-ja-en",  # Jepang ke Inggris
+    "zh": "Helsinki-NLP/opus-mt-zh-en",  # Mandarin ke Inggris
+}
+
+# Inisialisasi pipeline untuk setiap model
+translators = {}
 try:
-    logger.info("Loading translation model...")
-    # Load model dan tokenizer M2M100
-    model_name = "facebook/m2m100_418M"
-    model = M2M100ForConditionalGeneration.from_pretrained(model_name)
-    tokenizer = M2M100Tokenizer.from_pretrained(model_name)
-    logger.info("Model loaded successfully")
+    for lang, model_name in MODEL_MAP.items():
+        logger.info(f"Loading model for {lang}...")
+        translators[lang] = pipeline("translation", model=model_name)
+        logger.info(f"Model for {lang} loaded successfully")
 except Exception as e:
-    logger.error(f"Failed to load model: {str(e)}")
+    logger.error(f"Model initialization failed: {str(e)}")
     raise Exception(f"Model initialization failed: {str(e)}")
 
 def detect_language(text: str) -> str:
     """Deteksi bahasa dari teks menggunakan langdetect."""
     try:
         lang = langdetect.detect(text)
-        # Konversi kode bahasa ke format M2M100
-        lang_map = {
-            "id": "id",  # Indonesia
-            "fr": "fr",  # Prancis
-            "es": "es",  # Spanyol
-            "de": "de",  # Jerman
-            "ja": "ja",  # Jepang
-            "zh": "zh",  # Mandarin
-            "ru": "ru",  # Rusia
-            "th": "th",  # Thailand
-        }
-        return lang_map.get(lang, "en")  # Default ke Inggris jika tidak dikenali
+        return lang if lang in MODEL_MAP else "en"  # Default ke Inggris jika tidak didukung
     except Exception as e:
         logger.warning(f"Language detection failed: {str(e)}, defaulting to English")
         return "en"
 
-def translate_to_english(text: str, source_lang: str) -> str:
-    """Terjemahkan teks ke Inggris menggunakan M2M100."""
-    try:
-        # Set bahasa sumber
-        tokenizer.src_lang = source_lang
-        # Encode teks
-        encoded = tokenizer(text, return_tensors="pt")
-        # Generate terjemahan (ke Inggris)
-        generated_tokens = model.generate(**encoded, forced_bos_token_id=tokenizer.get_lang_id("en"))
-        # Decode hasil
-        translated_text = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
-        return translated_text
-    except Exception as e:
-        logger.error(f"Translation failed: {str(e)}")
-        raise HTTPException(status_code=500, detail=f"Translation failed: {str(e)}")
-
 @app.post("/translate")
 async def translate(text: str):
     if not text:
@@ -68,12 +51,15 @@ async def translate(text: str):
         # Deteksi bahasa
         source_lang = detect_language(text)
         logger.info(f"Detected source language: {source_lang}")
-        # Jika sudah Inggris, kembalikan teks asli
+        # Jika sudah Inggris atau bahasa tidak didukung, kembalikan teks asli
         if source_lang == "en":
             return {"translated_text": text}
         # Terjemahkan ke Inggris
-        translated_text = translate_to_english(text, source_lang)
-        return {"translated_text": translated_text}
+        translator = translators.get(source_lang)
+        if not translator:
+            raise HTTPException(status_code=400, detail=f"Translation not supported for language: {source_lang}")
+        result = translator(text)
+        return {"translated_text": result[0]["translation_text"]}
     except Exception as e:
         logger.error(f"Processing failed: {str(e)}")
         raise HTTPException(status_code=500, detail=f"Processing failed: {str(e)}")