Update main.py

main.py

@@ -1,86 +1,63 @@
 import os
-
-# cache dirs for HF
+# 2️⃣ Optional: force cache to writable directory
 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)
 
 
-from fastapi import FastAPI, Request, HTTPException
+from fastapi import FastAPI
 from pydantic import BaseModel
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
-import torch
-import secrets
-
+#from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
-# -----------------------------
-# 1️⃣ Generate / load API key
-# -----------------------------
-# Generate a new key once (uncomment to create)
-# print(secrets.token_hex(32))
+#from replacer import replace_words, replace_dict
+#from datasets import Dataset
+from transformers import MBartForConditionalGeneration, MBart50TokenizerFast, DataCollatorForSeq2Seq, Seq2SeqTrainer, Seq2SeqTrainingArguments
 
-# Or load from environment variable
-API_KEY = os.getenv("MY_API_KEY", "ec1a464f3948d7e9e0484efad4f71d0a0aa9f3fb37560697c42da0568b9fbac5")
+# -------------------------
+# 1️⃣ Get your HF token from Space Secrets
+# In your Space, go to Settings → Secrets → add HF_TOKEN
+#HF_TOKEN = os.environ.get("HF_TOKEN")
+#if HF_TOKEN is None:
+#    raise ValueError("HF_TOKEN not found. Please add it in your Space Secrets.")
 
-# -----------------------------
-# 2️⃣ Initialize FastAPI & model
-# -----------------------------
-app = FastAPI()
+# -------------------------
 
+# -------------------------
+# 3️⃣ Load private model
+model_name = "Gaoussin/bamalingua-bm_ml-fr_XX"
+model = MBartForConditionalGeneration.from_pretrained(model_name)
+tokenizer = MBart50TokenizerFast.from_pretrained("facebook/mbart-large-50")
 
-MODEL_NAME = "facebook/nllb-200-3.3B"
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, use_fast=False)
-model = AutoModelForSeq2SeqLM.from_pretrained(
-    MODEL_NAME, torch_dtype=torch.float16
-).to("cuda" if torch.cuda.is_available() else "cpu")
+tgt_lang = "bm_ml"
 
-# 🔑 Build lang_code_to_id manually
-lang_code_to_id = {
-    tok.strip("_"): idx
-    for tok, idx in zip(tokenizer.additional_special_tokens,
-                        tokenizer.additional_special_tokens_ids)
-}
+# -------------------------
+# 4️⃣ FastAPI app
+app = FastAPI()
 
 class TranslationRequest(BaseModel):
     text: str
-    src_lang: str
-    tgt_lang: str
-    
-
 
 @app.post("/translate")
-async def translate(req: TranslationRequest, request: Request):
-        # ---- Check API key ----
-    api_key = request.headers.get("X-API-KEY")
-    if api_key != API_KEY:
-        raise HTTPException(status_code=403, detail="Unauthorized API key")
-
-    # ---- Optional IP restriction ----
-    #client_ip = request.client.host
-    #if ALLOWED_IPS and client_ip not in ALLOWED_IPS:
-    #    raise HTTPException(status_code=403, detail="IP not allowed")
-
-    
-    # always set source language
-    tokenizer.src_lang = req.src_lang  # 👈 force the source language
+def translate(request: TranslationRequest):
+    #reverse_dict = {v: k for k, v in replace_dict.items()}
+    #text_for_ai = replace_words(request.text, reverse_dict)
     
     inputs = tokenizer(
-        req.text,
-        return_tensors="pt",
-    ).to(model.device)
-
-    # force target language
-    tgt_lang = req.tgt_lang
-    forced_bos_id = lang_code_to_id[tgt_lang]
+    request.text,
+    return_tensors="pt",
+    max_length=128,
+    truncation=True)
 
     outputs = model.generate(
-        **inputs,
-        forced_bos_token_id=forced_bos_id,
-        max_length=512#,
-        #num_beams=5,   # good for quality
-        #early_stopping=True
-    )
-
-    translation = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return {"translation": translation}
+    **inputs,
+    forced_bos_token_id=tokenizer.lang_code_to_id[tgt_lang])
+    text2 = tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]
+    
+    #text_for_user = replace_words(text2, replace_dict)
+    
+    return {"translation": text2[0].upper() + text2[1:]}
+            
+@app.get("/")
+def root():
+    return {"message": "API is running"}