Upload folder using huggingface_hub

app.py

@@ -7,12 +7,14 @@ from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
 from fastapi.responses import JSONResponse
 from huggingface_hub import HfApi
 from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+from peft import PeftModel
 from pydantic import BaseModel
 
 # ---------- Config ----------
-MODEL_ID = "wolethereader/STORM-OS-MT-3B"
+BASE_MODEL_ID = "wolethereader/STORM-OS-MT-3B"
+REVERSE_ADAPTER_ID = "wolethereader/STORM-OS-MT-3B-REVERSE"
 ORG_NAME = "wolethereader"
-TGT_LANG = "eng_Latn"
+EN = "eng_Latn"
 LANG_CODES = {"yo": "yor_Latn", "ha": "hau_Latn", "ig": "ibo_Latn", "pcm": "pcm_Latn"}
 VALID_LANGS = set(LANG_CODES.keys())
 MAX_TEXT_CHARS = 2000
@@ -69,27 +71,35 @@ def check_rate_limit(username: str):
 
 # ---------- Model ----------
 tokenizer = None
-model = None
+model = None  # PeftModel with the reverse adapter loaded but not necessarily active
 
 @app.on_event("startup")
 async def startup():
     global tokenizer, model
-    log_event("loading_model", model=MODEL_ID)
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
-    model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16)
+    log_event("loading_tokenizer_with_pcm_token", repo=REVERSE_ADAPTER_ID)
+    tokenizer = AutoTokenizer.from_pretrained(REVERSE_ADAPTER_ID)  # has pcm_Latn added
+
+    log_event("loading_base_model", repo=BASE_MODEL_ID)
+    base_model = AutoModelForSeq2SeqLM.from_pretrained(BASE_MODEL_ID, torch_dtype=torch.bfloat16)
+    base_model.resize_token_embeddings(len(tokenizer))
+
     device = "cuda" if torch.cuda.is_available() else "cpu"
-    model.to(device)
+    base_model.to(device)
+
+    log_event("attaching_reverse_adapter", repo=REVERSE_ADAPTER_ID)
+    model = PeftModel.from_pretrained(base_model, REVERSE_ADAPTER_ID, adapter_name="reverse")
     model.eval()
     log_event("model_loaded_ok", device=device)
 
 class TranslateRequest(BaseModel):
     text: str
-    source_lang: str
+    direction: str  # "forward" (local -> English) or "reverse" (English -> local)
+    lang: str        # the local language code, yo/ha/ig/pcm, regardless of direction
     max_new_tokens: int = 128
 
 @app.get("/")
 def root():
-    return {"status": "ok", "languages": sorted(VALID_LANGS), "engine": MODEL_ID}
+    return {"status": "ok", "languages": sorted(VALID_LANGS), "directions": ["forward", "reverse"], "engine": BASE_MODEL_ID}
 
 @app.get("/health")
 def health():
@@ -99,8 +109,10 @@ def health():
 async def translate(req: TranslateRequest, username: str = Depends(verify_org_token)):
     check_rate_limit(username)
 
-    if req.source_lang not in VALID_LANGS:
-        raise HTTPException(status_code=400, detail=f"source_lang must be one of {sorted(VALID_LANGS)}")
+    if req.lang not in VALID_LANGS:
+        raise HTTPException(status_code=400, detail=f"lang must be one of {sorted(VALID_LANGS)}")
+    if req.direction not in ("forward", "reverse"):
+        raise HTTPException(status_code=400, detail="direction must be 'forward' or 'reverse'")
     if not req.text or not req.text.strip():
         raise HTTPException(status_code=400, detail="text must not be empty")
     if len(req.text) > MAX_TEXT_CHARS:
@@ -109,26 +121,35 @@ async def translate(req: TranslateRequest, username: str = Depends(verify_org_to
     request_id = str(uuid.uuid4())
     start = time.time()
 
-    tokenizer.src_lang = LANG_CODES[req.source_lang]
+    if req.direction == "forward":
+        src_lang, tgt_lang = LANG_CODES[req.lang], EN
+        context = model.disable_adapter()
+    else:
+        src_lang, tgt_lang = EN, LANG_CODES[req.lang]
+        model.set_adapter("reverse")
+        context = None
+
+    tokenizer.src_lang = src_lang
     inputs = tokenizer(req.text, return_tensors="pt", truncation=True, max_length=128).to(model.device)
-    tgt_id = tokenizer.convert_tokens_to_ids(TGT_LANG)
+    tgt_id = tokenizer.convert_tokens_to_ids(tgt_lang)
+
     with torch.no_grad():
-        out = model.generate(
-            **inputs,
-            forced_bos_token_id=tgt_id,
-            max_new_tokens=req.max_new_tokens,
-            max_length=None,
-        )
+        if context is not None:
+            with context:
+                out = model.generate(**inputs, forced_bos_token_id=tgt_id, max_new_tokens=req.max_new_tokens, max_length=None)
+        else:
+            out = model.generate(**inputs, forced_bos_token_id=tgt_id, max_new_tokens=req.max_new_tokens, max_length=None)
+
     translated = tokenizer.decode(out[0], skip_special_tokens=True)
     elapsed_s = round(time.time() - start, 2)
 
     log_event("translate_ok", request_id=request_id, user=username,
-              source_lang=req.source_lang, elapsed_s=elapsed_s)
+              direction=req.direction, lang=req.lang, elapsed_s=elapsed_s)
 
     return {
         "request_id": request_id,
-        "source_lang": req.source_lang,
-        "target_lang": "en",
+        "direction": req.direction,
+        "lang": req.lang,
         "translated_text": translated,
         "elapsed_s": elapsed_s,
     }

requirements.txt

@@ -1,7 +1,9 @@
 fastapi
 uvicorn
 transformers==5.15.0
+peft==0.20.0
 torch
 sentencepiece
 accelerate
 python-multipart
+# add reverse toggle, force rebuild 1787637858