Update language.py

language.py

@@ -1,20 +1,9 @@
-"""
-language.py — small wrapper to expose a stable translation API for JusticeAI.
-
-Place this file next to your language.bin. The wrapper will try to load language.bin
-(via torch.load first, then via pickle) and adapt to common shapes:
-
-- object.translate(text, src, tgt) or object.translate_to_en(text, src)
-- object(text, src, tgt) if the loaded object is callable
-- dict-like mapping with keys (('src','tgt') -> function/string)
-- HF-like model object with .generate and an attached tokenizer attribute
-
-Functions exposed:
-- translate(text, src, tgt)
-- translate_to_en(text, src)
-- translate_from_en(text, tgt)
-"""
-
+# language.py — wrapper to expose a stable translation API for JusticeAI.
+# Tries to load language.bin (torch.load then pickle). Adapts common shapes and exposes:
+# - translate(text, src, tgt)
+# - translate_to_en(text, src)
+# - translate_from_en(text, tgt)
+# Also exposes model_info() for debugging/inspection.
 from pathlib import Path
 import logging
 import pickle
@@ -54,12 +43,23 @@ def _ensure_loaded():
     else:
         logger.info("language.bin not found in cwd")
 
-# Generic translator interface
+def model_info():
+    _ensure_loaded()
+    if _model is None:
+        return {"loaded": False}
+    info = {"loaded": True, "type": type(_model).__name__}
+    try:
+        info["repr"] = repr(_model)[:800]
+    except Exception:
+        info["repr"] = "<unreprable>"
+    info["has_translate"] = hasattr(_model, "translate")
+    info["has_translate_to_en"] = hasattr(_model, "translate_to_en")
+    info["has_translate_from_en"] = hasattr(_model, "translate_from_en")
+    info["callable"] = callable(_model)
+    info["dir"] = [n for n in dir(_model) if not n.startswith("_")]
+    return info
+
 def translate(text: str, src: str, tgt: str) -> str:
-    """
-    Translate text from src -> tgt using the loaded model.
-    If nothing useful is found, returns the original text.
-    """
     _ensure_loaded()
     if not text:
         return text
@@ -69,12 +69,11 @@ def translate(text: str, src: str, tgt: str) -> str:
     # 1) object has translate(text, src, tgt) or translate_to_en
     try:
         if hasattr(_model, "translate"):
-            # try both signatures (text, src, tgt) or (text, src_tgt) etc.
             try:
                 return _model.translate(text, src, tgt)
             except TypeError:
                 try:
-                    return _model.translate(text, src + "->" + tgt)
+                    return _model.translate(text, f"{src}->{tgt}")
                 except Exception:
                     pass
         if hasattr(_model, "translate_to_en") and tgt.lower() in ("en", "eng"):
@@ -113,7 +112,6 @@ def translate(text: str, src: str, tgt: str) -> str:
                     return fn(text)
                 if isinstance(fn, str):
                     return fn
-            # fallback to direct mapping by lang codes
             key2 = f"{src}->{tgt}"
             if key2 in _model:
                 val = _model[key2]
@@ -124,20 +122,15 @@ def translate(text: str, src: str, tgt: str) -> str:
     except Exception as e:
         logger.debug(f"dict-like model attempt failed: {e}")
 
-    # 4) HF-style object: has .generate and maybe a tokenizer at _model.tokenizer
+    # 4) HF-like object: has .generate and maybe a tokenizer at _model.tokenizer
     try:
-        tokenizer = None
         m = _model
-        if hasattr(m, "generate") and hasattr(m, "to") or hasattr(m, "device"):
-            # try to locate a tokenizer attached to model
-            tokenizer = getattr(m, "tokenizer", None)
-            if tokenizer is None and hasattr(m, "config") and hasattr(m, "config", "tokenizer_class"):
-                tokenizer = getattr(m, "tokenizer", None)
-            if tokenizer:
-                inputs = tokenizer([text], return_tensors="pt", truncation=True)
-                outputs = m.generate(**inputs, max_length=1024)
-                decoded = tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]
-                return decoded
+        tokenizer = getattr(m, "tokenizer", None)
+        if tokenizer and hasattr(m, "generate"):
+            inputs = tokenizer([text], return_tensors="pt", truncation=True)
+            outputs = m.generate(**inputs, max_length=1024)
+            decoded = tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]
+            return decoded
     except Exception as e:
         logger.debug(f"HF-like model attempt failed: {e}")
 
@@ -147,12 +140,10 @@ def translate(text: str, src: str, tgt: str) -> str:
 def translate_to_en(text: str, src: str) -> str:
     if not text:
         return text
-    # prefer an explicit helper if model provides it
     _ensure_loaded()
-    if _model is not None:
+    if _model is not None and hasattr(_model, "translate_to_en"):
         try:
-            if hasattr(_model, "translate_to_en"):
-                return _model.translate_to_en(text, src)
+            return _model.translate_to_en(text, src)
         except Exception:
             pass
     return translate(text, src, "en")
@@ -161,22 +152,32 @@ def translate_from_en(text: str, tgt: str) -> str:
     if not text:
         return text
     _ensure_loaded()
-    if _model is not None:
+    if _model is not None and hasattr(_model, "translate_from_en"):
         try:
-            if hasattr(_model, "translate_from_en"):
-                return _model.translate_from_en(text, tgt)
+            return _model.translate_from_en(text, tgt)
         except Exception:
             pass
     return translate(text, "en", tgt)
 
-# Simple CLI test helpers (run `python language.py` to exercise)
+# Optional: expose a detect function if the model has one, else None
+def detect_language(text: str) -> str:
+    _ensure_loaded()
+    if _model is None:
+        return None
+    for candidate in ("detect", "detect_language", "lang", "language"):
+        if hasattr(_model, candidate):
+            try:
+                return getattr(_model, candidate)(text)
+            except Exception:
+                pass
+    return None
+
 if __name__ == "__main__":
+    # simple CLI debug
     import sys
     _ensure_loaded()
+    print("model_info:", model_info())
     if len(sys.argv) >= 4:
-        _, src, tgt, txt = sys.argv[0], sys.argv[1], sys.argv[2], " ".join(sys.argv[3:])
-        out = translate(txt, src, tgt)
-        print(out)
-    else:
-        print("Usage: python language.py <src> <tgt> <text...>")
-        print("Example: python language.py es en 'hola mundo'")
+        _, src, tgt, *txt = sys.argv
+        txt = " ".join(txt)
+        print("translate:", translate(txt, src, tgt))