Fix: Corrected model loading on startup in FastAPI app

.dockerignore

@@ -0,0 +1,6 @@
+__pycache__/
+*.pyc
+.git/
+models/
+*.safetensors
+*.bin

Dockerfile

@@ -0,0 +1,17 @@
+FROM python:3.10-slim
+
+# Set working directory
+WORKDIR /app
+
+# Copy files
+COPY . .
+
+# Install dependencies
+RUN pip install --no-cache-dir --upgrade pip && \
+    pip install --no-cache-dir -r requirements.txt
+
+# Expose the default port used by Hugging Face Docker Spaces
+EXPOSE 7860
+
+# Start the FastAPI app using Uvicorn
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]

app.py

@@ -0,0 +1,9 @@
+# This imports the FastAPI app
+from translation_api import app
+
+# If needed, you can add any root route
+@app.get("/")
+def read_root():
+    return {"message": "Welcome to ChikaMo Translator API"}
+
+

requirements.txt

@@ -0,0 +1,16 @@
+transformers==4.41.2
+datasets==2.19.1
+torch>=2.2.0
+fastapi==0.111.0
+uvicorn==0.30.0
+huggingface_hub
+
+# googletrans==4.0.0-rc1  
+# ❌ Conflict
+deep-translator==1.11.4   
+#✅ Modern replacement
+
+python-dotenv==1.0.1
+tqdm==4.66.4
+scikit-learn==1.4.2
+rich==13.7.

translation_api.py

@@ -0,0 +1,161 @@
+# translate_api.py
+# Defines the FastAPI application for the Chiakamo Translator.
+
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+from transformers import MarianMTModel, MarianTokenizer, pipeline
+import torch
+import os # Import os for environment variables (e.g., HF_HOME)
+
+# Initialize the FastAPI application
+app = FastAPI(title="ChikaMo Translator API",
+              description="API for translating between Tagalog and English, with fallback to Helsinki-NLP models.")
+
+# --- Model Configuration ---
+# Define the Hugging Face model repository IDs for your custom models.
+# These should be your repositories on Hugging Face Hub (e.g., chikamov1/opus-mt-tl-en-chikamo)
+LOCAL_MODELS = {
+    "tl-en": "chikamov1/opus-mt-tl-en-chikamo", # Your custom Tagalog-to-English model
+    "en-tl": "chikamov1/opus-mt-en-tl-chikamo", # Your custom English-to-Tagalog model
+}
+
+# Define HuggingFace fallback models (Helsinki-NLP is a good choice)
+FALLBACK_MODELS = {
+    "tl-en": "Helsinki-NLP/opus-mt-tl-en",
+    "en-tl": "Helsinki-NLP/opus-mt-en-tl",
+}
+
+# Dictionaries to store loaded models and tokenizers to avoid reloading on every request
+loaded_models = {}
+loaded_tokenizers = {}
+fallback_pipelines = {}
+
+# --- Model Loading Functions ---
+
+# Function to get custom model and tokenizer (loads and caches)
+def get_model_and_tokenizer(pair: str):
+    """
+    Loads and caches the custom MarianMT model and tokenizer for a given language pair.
+    """
+    if pair in LOCAL_MODELS:
+        if pair not in loaded_models:
+            print(f"Attempting to load local model: {LOCAL_MODELS[pair]}")
+            try:
+                # Load model and tokenizer from your Hugging Face Hub repository
+                model = MarianMTModel.from_pretrained(LOCAL_MODELS[pair])
+                tokenizer = MarianTokenizer.from_pretrained(LOCAL_MODELS[pair])
+                loaded_models[pair] = model
+                loaded_tokenizers[pair] = tokenizer
+                print(f"Successfully loaded local model: {LOCAL_MODELS[pair]}")
+            except Exception as e:
+                print(f"Failed to load local model {LOCAL_MODELS[pair]}: {e}")
+                return None, None # Return None if loading fails
+        return loaded_models.get(pair), loaded_tokenizers.get(pair)
+    return None, None
+
+# Function to get fallback pipeline (loads and caches)
+def get_fallback_pipeline(pair: str):
+    """
+    Loads and caches a fallback translation pipeline for a given language pair.
+    """
+    if pair not in fallback_pipelines:
+        print(f"Attempting to load fallback model: {FALLBACK_MODELS[pair]}")
+        try:
+            # Use the pipeline abstraction for fallback models
+            pipe = pipeline("translation", model=FALLBACK_MODELS[pair])
+            fallback_pipelines[pair] = pipe
+            print(f"Successfully loaded fallback model: {FALLBACK_MODELS[pair]}")
+        except Exception as e:
+            print(f"Failed to load fallback model {FALLBACK_MODELS[pair]}: {e}")
+            return None # Return None if loading fails
+    return fallback_pipelines.get(pair)
+
+# --- Pydantic Model for Request Body ---
+class TranslationRequest(BaseModel):
+    """
+    Defines the structure of the incoming JSON request for translation.
+    """
+    source_lang: str # e.g., "tl" for Tagalog, "en" for English
+    target_lang: str # e.g., "en" for English, "tl" for Tagalog
+    text: str        # The text to be translated
+
+# --- API Endpoints ---
+
+# Root endpoint for basic API information
+@app.get("/")
+def read_root():
+    """
+    Returns a welcome message and API status.
+    """
+    return {"message": "Welcome to ChikaMo Translator API", "status": "running"}
+
+# Translation endpoint
+@app.post("/translate")
+def translate_text_endpoint(req: TranslationRequest):
+    """
+    Translates text between specified source and target languages.
+    Prioritizes custom models, falls back to Helsinki-NLP if custom fails or is not found.
+    """
+    pair = f"{req.source_lang.lower()}-{req.target_lang.lower()}"
+    text = req.text.strip()
+
+    if not text:
+        raise HTTPException(status_code=400, detail="Input text is empty.")
+
+    translated_output = ""
+    fallback_used = False
+
+    # Try to use local (custom) model first
+    model, tokenizer = get_model_and_tokenizer(pair)
+    if model and tokenizer:
+        try:
+            inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
+            # Ensure model is on CPU if no GPU is available, or move to GPU if present
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+            model.to(device)
+            inputs = {k: v.to(device) for k, v in inputs.items()}
+
+            with torch.no_grad():
+                translated = model.generate(**inputs)
+            translated_output = tokenizer.decode(translated[0], skip_special_tokens=True)
+            fallback_used = False
+        except Exception as e:
+            print(f"Error during custom model translation for pair {pair}: {e}. Attempting fallback.")
+            # Clear to force fallback path
+            model, tokenizer = None, None
+    
+    # If custom model failed or wasn't available, try fallback
+    if not model and pair in FALLBACK_MODELS:
+        pipe = get_fallback_pipeline(pair)
+        if pipe:
+            try:
+                translated_output = pipe(text)[0]["translation_text"]
+                fallback_used = True
+            except Exception as e:
+                print(f"Error during fallback model translation for pair {pair}: {e}.")
+                raise HTTPException(status_code=500, detail=f"Translation failed for pair {pair} with both custom and fallback models.")
+        else:
+            raise HTTPException(status_code=500, detail=f"Fallback model for pair {pair} could not be loaded.")
+    elif not model: # No custom model and no fallback defined
+        raise HTTPException(status_code=400, detail=f"Unsupported language pair: {pair}, and no fallback model configured.")
+
+
+    return {
+        "translation": translated_output,
+        "source_lang": req.source_lang,
+        "target_lang": req.target_lang,
+        "fallback_used": fallback_used
+    }
+
+# --- Application Startup Event ---
+# This ensures models are loaded when the FastAPI app starts up.
+# This is crucial for Hugging Face Spaces where the app is started by Uvicorn.
+@app.on_event("startup")
+async def startup_event():
+    # Pre-load all local and fallback models during startup
+    print("Pre-loading models during application startup...")
+    for pair in LOCAL_MODELS.keys():
+        get_model_and_tokenizer(pair) # Attempt to load custom models
+    for pair in FALLBACK_MODELS.keys():
+        get_fallback_pipeline(pair) # Attempt to load fallback pipelines
+    print("Model pre-loading complete.")