Update app.py

app.py

@@ -5,55 +5,241 @@ import gradio as gr
 import os
 import asyncio
 import atexit
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
 
-# Constants
-MAX_INPUT_LENGTH = 10000  # Maximum characters allowed
+# ==================== Constants ====================
+MAX_INPUT_LENGTH = 10000  # OpenLID character limit
+COMMONLINGUA_MAX_BYTES = 512  # CommonLingua byte limit
 
-# Preprocessing patterns
+# ==================== OpenLID Setup ====================
+print("Loading OpenLID-v3 model...")
+openlid_path = hf_hub_download(
+    repo_id="HPLT/OpenLID-v3",
+    filename="openid-v3.bin"
+)
+openlid_model = fasttext.load_model(openlid_path)
+print("OpenLID-v3 loaded successfully!")
+
+# Preprocessing patterns for OpenLID
 NONWORD_REPLACE_STR = r"[^\p{Word}\p{Zs}]|\d"
 NONWORD_REPLACE_PATTERN = regex.compile(NONWORD_REPLACE_STR)
 SPACE_PATTERN = regex.compile(r"\s\s+")
 
-def preprocess(text):
-    """Preprocess text for language identification."""
+def openlid_preprocess(text):
+    """Preprocess text for OpenLID-v3."""
     text = text.strip().replace('\n', ' ').lower()
     text = regex.sub(SPACE_PATTERN, " ", text)
     text = regex.sub(NONWORD_REPLACE_PATTERN, "", text)
     return text
 
-# Load model once at startup
-print("Loading OpenLID-v3 model...")
-model_path = hf_hub_download(
-    repo_id="HPLT/OpenLID-v3", 
-    filename="openlid-v3.bin"
+# ==================== CommonLingua Setup ====================
+# Inline model architecture (from model.py) so no extra file is needed
+class ByteNgramEmbed(nn.Module):
+    def __init__(self, num_buckets=4096, embed_dim=64, n=3):
+        super().__init__()
+        self.n = n
+        self.num_buckets = num_buckets
+        self.embed = nn.Embedding(num_buckets, embed_dim)
+
+    def forward(self, byte_ids):
+        B, T = byte_ids.shape
+        clamped = byte_ids.clamp(max=255)
+        padded = F.pad(clamped, (0, self.n - 1), value=0)
+        h = torch.zeros(B, T, dtype=torch.long, device=byte_ids.device)
+        for i in range(self.n):
+            h = h * 257 + padded[:, i:i + T]
+        return self.embed(h % self.num_buckets)
+
+
+class ByteConvBlock(nn.Module):
+    def __init__(self, d_model, kernel_size=15, expand=2):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(d_model)
+        self.pad = kernel_size - 1
+        self.conv = nn.Conv1d(d_model, d_model, kernel_size, groups=d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+        ffn = d_model * expand
+        self.ffn_gate = nn.Linear(d_model, ffn, bias=False)
+        self.ffn_up = nn.Linear(d_model, ffn, bias=False)
+        self.ffn_down = nn.Linear(ffn, d_model, bias=False)
+
+    def forward(self, x):
+        residual = x
+        x = self.norm1(x).transpose(1, 2)
+        x = F.pad(x, (self.pad, 0))
+        x = F.silu(self.conv(x)).transpose(1, 2)
+        x = residual + x
+
+        residual = x
+        x = self.norm2(x)
+        x = self.ffn_down(F.silu(self.ffn_gate(x)) * self.ffn_up(x))
+        return residual + x
+
+
+def _rope(q, k):
+    head_dim = q.shape[-1]
+    seq_len = q.shape[-2]
+    freqs = 1.0 / (10000.0 ** (torch.arange(0, head_dim, 2, device=q.device).float() / head_dim))
+    t = torch.arange(seq_len, device=q.device)
+    a = torch.outer(t, freqs)
+    cos = a.cos().to(q.dtype)
+    sin = a.sin().to(q.dtype)
+
+    def rot(x):
+        x1, x2 = x[..., : head_dim // 2], x[..., head_dim // 2:]
+        return torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin], dim=-1)
+
+    return rot(q), rot(k)
+
+
+class ByteAttnBlock(nn.Module):
+    def __init__(self, d_model, n_heads=4, expand=2):
+        super().__init__()
+        self.n_heads = n_heads
+        self.head_dim = d_model // n_heads
+        self.norm1 = nn.LayerNorm(d_model)
+        self.qkv = nn.Linear(d_model, 3 * d_model, bias=False)
+        self.out_proj = nn.Linear(d_model, d_model, bias=False)
+        self.norm2 = nn.LayerNorm(d_model)
+        ffn = d_model * expand
+        self.ffn_gate = nn.Linear(d_model, ffn, bias=False)
+        self.ffn_up = nn.Linear(d_model, ffn, bias=False)
+        self.ffn_down = nn.Linear(ffn, d_model, bias=False)
+
+    def forward(self, x):
+        B, T, D = x.shape
+        residual = x
+        h = self.norm1(x)
+        qkv = self.qkv(h).reshape(B, T, 3, self.n_heads, self.head_dim)
+        q, k, v = (t.transpose(1, 2) for t in qkv.unbind(dim=2))
+        q, k = _rope(q, k)
+        attn = (q @ k.transpose(-2, -1)) / (self.head_dim ** 0.5)
+        attn = attn.softmax(dim=-1)
+        out = (attn @ v).transpose(1, 2).contiguous().view(B, T, D)
+        x = residual + self.out_proj(out)
+
+        residual = x
+        h = self.norm2(x)
+        h = self.ffn_down(F.silu(self.ffn_gate(h)) * self.ffn_up(h))
+        return residual + h
+
+
+class ByteHybrid(nn.Module):
+    def __init__(
+        self,
+        num_classes,
+        d_model=256,
+        n_conv=3,
+        n_attn=1,
+        n_heads=4,
+        ffn_expand=2,
+        max_len=512,
+        conv_kernel=15,
+        ngram_buckets=0,
+        ngram_dim=64,
+    ):
+        super().__init__()
+        self.max_len = max_len
+        self.embed = nn.Embedding(257, d_model, padding_idx=256)
+
+        self.ngram_embed = None
+        if ngram_buckets > 0:
+            self.ngram_embed = ByteNgramEmbed(ngram_buckets, ngram_dim, n=3)
+            self.ngram_proj = nn.Linear(ngram_dim, d_model, bias=False)
+
+        self.conv_layers = nn.ModuleList(
+            [ByteConvBlock(d_model, conv_kernel, ffn_expand) for _ in range(n_conv)]
+        )
+        self.attn_layers = nn.ModuleList(
+            [ByteAttnBlock(d_model, n_heads, ffn_expand) for _ in range(n_attn)]
+        )
+        self.final_norm = nn.LayerNorm(d_model)
+        self.head = nn.Sequential(
+            nn.Linear(d_model, d_model),
+            nn.GELU(),
+            nn.Dropout(0.1),
+            nn.Linear(d_model, num_classes),
+        )
+
+    def forward(self, byte_ids):
+        pad_mask = byte_ids != 256
+        x = self.embed(byte_ids)
+        if self.ngram_embed is not None:
+            x = x + self.ngram_proj(self.ngram_embed(byte_ids))
+        for layer in self.conv_layers:
+            x = layer(x)
+        for layer in self.attn_layers:
+            x = layer(x)
+        x = self.final_norm(x)
+        mask = pad_mask.unsqueeze(-1).to(x.dtype)
+        x = (x * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1)
+        return self.head(x)
+
+
+CONFIGS = {
+    "base_ngram": dict(
+        d_model=256, n_conv=3, n_attn=1, n_heads=4, conv_kernel=15,
+        ngram_buckets=4096, ngram_dim=64,
+    ),
+}
+
+def commonlingua_encode(texts, max_len):
+    out = np.full((len(texts), max_len), 256, dtype=np.int64)
+    for i, t in enumerate(texts):
+        if not isinstance(t, str):
+            t = "" if t is None else str(t)
+        raw = t.encode("utf-8", errors="replace")[:max_len]
+        if raw:
+            out[i, :len(raw)] = np.frombuffer(raw, dtype=np.uint8)
+    return torch.from_numpy(out)
+
+
+@torch.no_grad()
+def commonlingua_predict(model, texts, idx2lang, max_len, device, top_k=3):
+    """Returns a list of [(lang, prob), ...] (one list per text, top-k entries each)."""
+    out = []
+    batch = commonlingua_encode(texts, max_len).to(device)
+    probs = torch.softmax(model(batch).float(), dim=-1)
+    top_p, top_idx = probs.topk(top_k, dim=-1)
+    for p_row, idx_row in zip(top_p.cpu().tolist(), top_idx.cpu().tolist()):
+        out.append([(idx2lang[j], float(p)) for p, j in zip(p_row, idx_row)])
+    return out
+
+
+print("Loading CommonLingua model...")
+commonlingua_path = hf_hub_download(
+    repo_id="PleIAs/CommonLingua",
+    filename="model.pt"
 )
-model = fasttext.load_model(model_path)
-print("Model loaded successfully!")
+ckpt = torch.load(commonlingua_path, map_location="cpu", weights_only=False)
+commonlingua_model = ByteHybrid(
+    num_classes=ckpt["num_classes"],
+    max_len=ckpt["max_len"],
+    **CONFIGS[ckpt["config"]]
+)
+commonlingua_model.load_state_dict(ckpt["model_state_dict"])
+commonlingua_model.eval()
 
-def predict_language(text, top_k=3, threshold=0.5):
-    """
-    Predict language of input text.
-    
-    Args:
-        text: Input text to analyze
-        top_k: Number of top predictions to return (1-10)
-        threshold: Confidence threshold (0.0-1.0)
-    """
-    # Check input length first
-    if len(text) > MAX_INPUT_LENGTH:
-        return f"**Error**: Input too long ({len(text):,} characters). Maximum allowed is {MAX_INPUT_LENGTH:,} characters."
-    
+device = "cuda" if torch.cuda.is_available() else "cpu"
+commonlingua_model = commonlingua_model.to(device)
+commonlingua_idx2lang = {v: k for k, v in ckpt["lang2idx"].items()}
+commonlingua_max_len = ckpt["max_len"]
+print(f"CommonLingua loaded successfully! ({len(commonlingua_idx2lang)} languages, device={device})")
+
+# ==================== Prediction Functions ====================
+def predict_openlid(text, top_k=3, threshold=0.5):
+    """Predict language using OpenLID-v3."""
     if not text or not text.strip():
         return "Please enter some text to analyze."
     
-    # Preprocess
-    processed_text = preprocess(text)
-    
+    processed_text = openlid_preprocess(text)
     if not processed_text.strip():
         return "Text contains no valid characters for language identification."
     
-    # Get predictions
-    predictions = model.predict(
+    predictions = openlid_model.predict(
         text=processed_text,
         k=min(top_k, 10),
         threshold=threshold,
@@ -61,18 +247,58 @@ def predict_language(text, top_k=3, threshold=0.5):
     )
     
     labels, scores = predictions
-    
-    # Format results
     results = []
     for label, score in zip(labels, scores):
-        # Remove __label__ prefix and format
         lang_code = label.replace("__label__", "")
         confidence = float(score) * 100
         results.append(f"**{lang_code}**: {confidence:.2f}%")
     
-    return "\n\n".join(results)
+    return "\n\n".join(results) if results else "No predictions above threshold."
+
+
+def predict_commonlingua(text, top_k=3):
+    """Predict language using CommonLingua."""
+    if not text or not text.strip():
+        return "Please enter some text to analyze."
+    
+    results = commonlingua_predict(
+        commonlingua_model, [text], commonlingua_idx2lang,
+        commonlingua_max_len, device, top_k=min(top_k, 10)
+    )
+    
+    formatted = []
+    for lang, prob in results[0]:
+        formatted.append(f"**{lang}**: {prob*100:.2f}%")
+    return "\n\n".join(formatted)
+
 
-# Cleanup function to prevent async errors on shutdown
+def predict_both(text, top_k=3, threshold=0.5):
+    """
+    Run both models and return combined results.
+    Returns tuple: (openlid_result, commonlingua_result, status_message)
+    """
+    # Check OpenLID length limit
+    if len(text) > MAX_INPUT_LENGTH:
+        return (
+            f"**Error**: Input too long ({len(text):,} characters). Maximum allowed is {MAX_INPUT_LENGTH:,} characters.",
+            f"**Error**: Input too long ({len(text):,} characters). Maximum allowed is {MAX_INPUT_LENGTH:,} characters.",
+            "❌ Input exceeds maximum length."
+        )
+    
+    # Check CommonLingua byte limit
+    byte_length = len(text.encode('utf-8'))
+    if byte_length > COMMONLINGUA_MAX_BYTES:
+        status = f"⚠️ Warning: Input is {byte_length} bytes. CommonLingua works best with ≤{COMMONLINGUA_MAX_BYTES} bytes (first {COMMONLINGUA_MAX_BYTES} bytes will be used)."
+    else:
+        status = f"✅ Input length: {len(text):,} chars | {byte_length} bytes"
+    
+    openlid_result = predict_openlid(text, top_k, threshold)
+    commonlingua_result = predict_commonlingua(text, top_k)
+    
+    return openlid_result, commonlingua_result, status
+
+
+# ==================== Cleanup ====================
 def cleanup():
     try:
         loop = asyncio.get_event_loop()
@@ -85,14 +311,15 @@ def cleanup():
 
 atexit.register(cleanup)
 
-# Create Gradio interface
-with gr.Blocks(title="OpenLID-v3 Language Identification") as demo:
-    # Use HTML with target="_blank" to open in new tab
+# ==================== Gradio Interface ====================
+with gr.Blocks(title="OpenLID-v3 vs CommonLingua") as demo:
     gr.HTML("""
-    <h1>OpenLID-v3 Language Identifier</h1>
-    <p>Identify the language of any text with state-of-the-art accuracy.<br>
-    Supports 194+ language varieties.</p>
-    <p><em>Model: <a href="https://huggingface.co/HPLT/OpenLID-v3" target="_blank" rel="noopener noreferrer">HPLT/OpenLID-v3</a></em></p>
+    <h1>🔍 Language Identification: OpenLID-v3 vs CommonLingua</h1>
+    <p>Compare two state-of-the-art language identification models side-by-side.</p>
+    <p>
+      <em>OpenLID-v3</em>: <a href="https://huggingface.co/HPLT/OpenLID-v3" target="_blank">HPLT/OpenLID-v3</a> (fastText, 194+ languages)<br>
+      <em>CommonLingua</em>: <a href="https://huggingface.co/PleIAs/CommonLingua" target="_blank">PleIAs/CommonLingua</a> (byte-level CNN+Attention, 334 languages, 2.35M params)
+    </p>
     """)
     
     with gr.Row():
@@ -102,29 +329,27 @@ with gr.Blocks(title="OpenLID-v3 Language Identification") as demo:
                 placeholder="Enter text to identify its language...",
                 lines=5,
                 max_lines=10,
-                max_length=MAX_INPUT_LENGTH  # Also enforce in UI
+                max_length=MAX_INPUT_LENGTH
             )
             with gr.Row():
                 top_k = gr.Slider(
-                    minimum=1, 
-                    maximum=10, 
-                    value=3, 
-                    step=1, 
+                    minimum=1, maximum=10, value=3, step=1,
                     label="Top-K Predictions"
                 )
                 threshold = gr.Slider(
-                    minimum=0.0, 
-                    maximum=1.0, 
-                    value=0.5, 
-                    step=0.05, 
-                    label="Confidence Threshold"
+                    minimum=0.0, maximum=1.0, value=0.5, step=0.05,
+                    label="OpenLID Confidence Threshold"
                 )
-            submit_btn = gr.Button("Identify Language", variant="primary")
+            submit_btn = gr.Button("🔍 Identify Language", variant="primary")
+            status = gr.Textbox(label="Status", interactive=False)
         
+    with gr.Row():
+        with gr.Column():
+            openlid_output = gr.Markdown(label="OpenLID-v3 Predictions")
         with gr.Column():
-            output = gr.Markdown(label="Predictions")
+            commonlingua_output = gr.Markdown(label="CommonLingua Predictions")
     
-    # Examples with Kabyle and Occitan as defaults
+    # Examples
     gr.Examples(
         examples=[
             ["Asebter-a yura s wudem awurman d amagrad s tutlayt taqbaylit."],
@@ -135,42 +360,42 @@ with gr.Blocks(title="OpenLID-v3 Language Identification") as demo:
             ["El rápido zorro marrón salta sobre el perro perezoso."],
             ["Быстрая коричневая лисица прыгает через ленивую собаку."],
             ["快速的棕色狐狸跳过了懒惰的狗。"],
+            ["Wikipédia est une encyclopédie universelle, multilingue."],
+            ["CommonLingua est un modèle d'identification de langue très léger."],
         ],
         inputs=input_text,
-        label="Try these examples (Kabyle and Occitan featured)"
+        label="Try these examples"
     )
     
     gr.Markdown(f"""
     ### Tips for best results:
-    - Text is automatically preprocessed (lowercased, normalized)
-    - Longer texts generally give more accurate predictions
-    - The model supports 194+ language varieties
-    - Use higher thresholds to filter out uncertain predictions
-    - **Maximum input length: {MAX_INPUT_LENGTH:,} characters**
+    - **OpenLID-v3**: Text is automatically preprocessed (lowercased, normalized). Longer texts generally give more accurate predictions. Max {MAX_INPUT_LENGTH:,} characters.
+    - **CommonLingua**: Operates directly on raw UTF-8 bytes (no tokenizer). Designed for paragraph-level corpus curation. Works best with ≤{COMMONLINGUA_MAX_BYTES} bytes. Not assessed on very short segments.
+    - Use the **Top-K** slider to see more alternative predictions.
+    - Use the **Threshold** slider to filter out uncertain OpenLID predictions (does not affect CommonLingua).
     """)
     
     # Event handlers
     submit_btn.click(
-        fn=predict_language,
+        fn=predict_both,
         inputs=[input_text, top_k, threshold],
-        outputs=output
+        outputs=[openlid_output, commonlingua_output, status]
     )
     
     input_text.submit(
-        fn=predict_language,
+        fn=predict_both,
         inputs=[input_text, top_k, threshold],
-        outputs=output
+        outputs=[openlid_output, commonlingua_output, status]
     )
 
 if __name__ == "__main__":
-    # Get port from environment (HF Spaces sets this)
     port = int(os.environ.get("PORT", 7860))
     
     try:
         demo.launch(
             server_name="0.0.0.0",
             server_port=port,
-            ssr_mode=False,  # Disable experimental SSR to prevent the error
+            ssr_mode=False,
             share=False,
             show_error=True
         )