Update main.py

main.py

@@ -1,50 +1,98 @@
 import os
 import torch
-from fastapi import FastAPI
+from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
-from transformers import MBartForConditionalGeneration, MBart50Tokenizer
+# Note: Keep the imports together for clarity
+from transformers import NllbTokenizer, AutoModelForSeq2SeqLM, Seq2SeqTrainer, Seq2SeqTrainingArguments, DataCollatorForSeq2Seq
 
-# 1️⃣ Cache (optional)
+# =====================
+# 1️⃣ Environment / Cache
+# =====================
+# Setting cache environment variables for Hugging Face
 os.environ["HF_HOME"] = "/tmp/hf"
 os.environ["TRANSFORMERS_CACHE"] = "/tmp/hf"
 os.environ["HF_DATASETS_CACHE"] = "/tmp/hf"
 os.makedirs("/tmp/hf", exist_ok=True)
 
-# 2️⃣ HF TOKEN
-HF_TOKEN = os.environ.get("mySpace")
-if HF_TOKEN is None:
-    raise ValueError("HF_TOKEN not found. Please add it in your Space Secrets.")
-
-# 3️⃣ DEVICE
+# =====================
+# 2️⃣ Device
+# =====================
 device = "cuda" if torch.cuda.is_available() else "cpu"
+print(f"Using device: {device}")
 
-model_name = "Gaoussin/bamalingua-bm_ml-fr_XX"
-tokenizer = MBart50Tokenizer.from_pretrained(model_name)
-model = MBartForConditionalGeneration.from_pretrained("Gaoussin/bamalingua-bm_ml-fr_XX")
-#####
-
-
-# 5️⃣ Translation function
-def translateTo(text, src_lang, tgt_lang):
-    tokenizer.src_lang = src_lang
-    inputs = tokenizer(text, return_tensors="pt").to(device)
-    tgt_id = tokenizer.lang_code_to_id[tgt_lang]
-    generated = model.generate(**inputs, forced_bos_token_id=tgt_id)
-    return tokenizer.decode(generated[0], skip_special_tokens=True)
+# =====================
+# 3️⃣ Load Model & Tokenizer
+# =====================
+# Charger le modèle et le tokenizer NLLB
+try:
+    model_name = "Gaoussin/bamalingua-2"
+    tokenizer = NllbTokenizer.from_pretrained(model_name)
+    # Move model to the selected device (CPU or GPU)
+    model = AutoModelForSeq2SeqLM.from_pretrained(model_name).to(device)
+    print(f"Model '{model_name}' loaded successfully on {device}.")
+except Exception as e:
+    print(f"Error loading model or tokenizer: {e}")
+    # In a real application, you might exit or handle this more gracefully
 
-# 6️⃣ FastAPI
+# =====================
+# 4️⃣ FastAPI setup - Define Input and Output Schemas
+# =====================
 app = FastAPI()
 
+# Input schema
 class TranslationRequest(BaseModel):
     text: str
-    src_lang: str
-    tgt_lang: str
+    src_lang: str  # e.g., "bam_Latn"
+    tgt_lang: str  # e.g., "fra_Latn"
+
+# Output schema (THE FIX: ensures both fields are returned)
+class TranslationResponse(BaseModel):
+    """
+    Ensures both the translated text and the app version ID are included 
+    in the response JSON.
+    """
+    translation: str
+    appVersionId: str
+
+# =====================
+# 5️⃣ Translation function - Restored to user's original logic
+# =====================
+def translateTo(text, src, tgt):
+    tokenizer.src_lang = src
+    tokenizer.tgt_lang = tgt
+    print(tokenizer.src_lang, tokenizer.tgt_lang)
+    
+    # Prepare input for the model
+    # We explicitly move the inputs to the same device as the model
+    inputs = tokenizer(text, return_tensors="pt").to(device)
+    
+    # Generate translation using the user's logic
+    output = model.generate(**inputs, max_length=128)
+    
+    # Decode the output
+    return tokenizer.decode(output[0], skip_special_tokens=True)
 
-@app.post("/translate")
+# =====================
+# 6️⃣ API Endpoints - Applying the Response Model
+# =====================
+@app.post("/translate", response_model=TranslationResponse) # <-- Fix remains here
 def translate(request: TranslationRequest):
-    output = translateTo(request.text, request.src_lang, request.tgt_lang)
-    return {"translation": output}
+    try:
+        result = translateTo(request.text, request.src_lang, request.tgt_lang)
+        appVersionId = "App Version id = 2"
+        
+        # Return the dictionary matching the TranslationResponse schema
+        return {"translation": result, "appVersionId": appVersionId}
+        
+    except Exception as e:
+        print(f"An error occurred during translation: {e}")
+        # When raising an HTTPException, the response model is bypassed,
+        # and a standard JSON error is returned.
+        raise HTTPException(
+            status_code=500, 
+            detail=f"Translation failed: {str(e)}"
+        )
 
 @app.get("/")
 def root():
-    return {"message": "API is running ✅"}
+    return {"message": "API is running 🚀"}