Upload 4 files

app.py

@@ -1,207 +1,194 @@
-import os
+import html
+import json
+import re
+import tempfile
 from functools import lru_cache
 from pathlib import Path
 
 import gradio as gr
-import pandas as pd
-import spaces
-import stanza
-from huggingface_hub import snapshot_download
-
-
-MODEL_REPO_ID = "usmannawaz/old-church-slavonic-tokenizer-lemmatizer"
-MODEL_VARIANT = "combined"
-
-LANG = "cu"
-POS_PACKAGE = "proiel_nocharlm"
-TOKENIZER_PACKAGE = "proiel"
-LEMMATIZER_PACKAGE = "proiel_nocharlm"
-
-HF_LOCAL_DIR = "hf_models/old-church-slavonic-tokenizer-lemmatizer"
-STANZA_DIR = os.path.join(os.path.expanduser("~"), "stanza_resources")
-
-os.environ["STANZA_RESOURCES_DIR"] = STANZA_DIR
-
-EXAMPLE_TEXT = "тѣма пророчьскꙑи ликъ д͞ховьнꙑѧ благодѣти исплъни сѧ"
-
-REPLACEMENTS1 = {
-    "": "и",
-    "": "ч",
-    "": "ѥ",
-    "": "н",
-    "": ":~",
-    "": "~",
-    "ⷣ": "ⷣ҇",
-    "": "̅",
-    "": "҆̀",
-    "": "ⷮ",
-    "": "ᲈ",
-    "": "Ч",
-    "": "ꙋ",
-    "": "чᲈ",
-    "": "҅́",
-    "": "꙾",
-    "": "ⷮ",
-    "": "҆́",
-    "ⷭⷭ": "҇",
-    "": "҆́",
-    "": "н҇",
-    "": "ꙶ",
-    "": "оу",
-    "": "ꙁ",
-    "": "ꙑ",
-    "": "ⱑ",
-    "": "с",
-    "": "Ѱ",
-    "": "҇",
-    "": "ⱉ",
-    "": "͡",
-    "": "҃",
-    "ⷮ": "ⷮ҇",
-    "ⷯ": "ⷯ҇",
-    "": "ꙩ́",
-    "ⷤ҆": "ⷤ҇",
-    "": "",
-    "": "҃",
-    "": "꙯",
-    "": "҃",
+import torch
+from huggingface_hub import hf_hub_download
+
+from inference import load_lemmatizer, load_registry, MODEL_REPO_ID, MODEL_ROOT
+
+
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+REGISTRY = load_registry("models_registry.json")
+
+TARGET_COL_IDX = 2
+BATCH_SIZE = 32
+
+
+
+
+def _display_name(item):
+    return f"{item['language']} - {item['treebank']}"
+
+
+LANGUAGES = sorted({item["language"] for item in REGISTRY.values()})
+
+DISPLAY_TO_ID = {
+    _display_name(item): model_id
+    for model_id, item in REGISTRY.items()
 }
 
-FONT_DIR = Path("fonts")
 
-if FONT_DIR.exists():
-    gr.set_static_paths(paths=[FONT_DIR])
+def treebank_choices(language):
+    choices = []
+
+    for model_id, item in REGISTRY.items():
+        if item["language"] == language:
+            choices.append(_display_name(item))
+
+    return sorted(choices)
+
+
+def default_language():
+    if "Old Church Slavonic" in LANGUAGES:
+        return "Old Church Slavonic"
+
+    return LANGUAGES[0] if LANGUAGES else None
+
+
+def default_treebank(language):
+    choices = treebank_choices(language)
+    preferred = "Old Church Slavonic - PROIEL"
+
+    if preferred in choices:
+        return preferred
+
+    return choices[0] if choices else None
+
 
 CUSTOM_CSS = """
-@font-face {
-    font-family: "Menaion";
-    src: url("/gradio_api/file=fonts/Menaion.otf") format("opentype");
-    font-style: normal;
-    font-weight: normal;
-    font-display: swap;
-}
 body {
-    background: linear-gradient(135deg, #f8fafc 0%, #eef6ff 50%, #f3f7fb 100%) !important;
+    background: linear-gradient(135deg, #eaf3ff 0%, #ffffff 48%, #dbeafe 100%);
 }
 .gradio-container {
     max-width: 980px !important;
     margin: auto !important;
-    font-family: Inter, Arial, sans-serif !important;
+    font-family: Arial, Helvetica, sans-serif !important;
 }
-#hero {
-    background: linear-gradient(135deg, #0f172a, #0f5f8f);
-    border-radius: 28px;
-    padding: 34px;
-    margin-bottom: 22px;
-    box-shadow: 0 22px 60px rgba(15, 23, 42, 0.22);
+#main-card {
+    background: #ffffff;
+    border: 1px solid #bfdbfe;
+    border-radius: 26px;
+    padding: 30px;
+    box-shadow: 0 20px 50px rgba(15, 23, 42, 0.16);
 }
-#hero h1 {
-    color: #ffffff;
-    font-size: 2.4rem;
-    margin: 0 0 10px 0;
+#title {
+    text-align: center;
+    color: #020617;
+    font-size: 2.5rem;
     font-weight: 900;
+    margin-bottom: 0.25rem;
 }
-#hero p {
-    color: #e0f2fe;
+#subtitle {
+    text-align: center;
+    color: #1e40af;
     font-size: 1.05rem;
     line-height: 1.55;
-    margin: 0;
+    margin-bottom: 1.6rem;
 }
-#main-card {
-    background: #ffffff;
-    border: 1px solid #cbd5e1;
-    border-radius: 24px;
-    padding: 24px;
-    box-shadow: 0 18px 45px rgba(15, 23, 42, 0.10);
+#badge-row {
+    text-align: center;
+    margin-bottom: 1.2rem;
 }
-#output-card {
-    background: #ffffff;
-    border: 1px solid #cbd5e1;
-    border-radius: 24px;
-    padding: 22px;
-    margin-top: 22px;
-    box-shadow: 0 18px 45px rgba(15, 23, 42, 0.10);
+#badge-row span {
+    display: inline-block;
+    background: #eff6ff;
+    color: #1e3a8a;
+    border: 1px solid #bfdbfe;
+    border-radius: 999px;
+    padding: 7px 13px;
+    margin: 4px;
+    font-size: 0.88rem;
+    font-weight: 700;
 }
-textarea, input {
+textarea, input, select {
     border-radius: 16px !important;
-    border: 1.5px solid #94a3b8 !important;
+    border: 1.5px solid #2563eb !important;
     background: #ffffff !important;
-    color: #0f172a !important;
+    color: #020617 !important;
+    box-shadow: 0 6px 16px rgba(37, 99, 235, 0.08) !important;
 }
 textarea:focus, input:focus {
-    border-color: #0f5f8f !important;
-    box-shadow: 0 0 0 3px rgba(14, 116, 144, 0.18) !important;
+    border-color: #1d4ed8 !important;
+    box-shadow: 0 0 0 3px rgba(37, 99, 235, 0.18) !important;
 }
 label {
-    color: #0f172a !important;
-    font-weight: 850 !important;
+    color: #020617 !important;
+    font-weight: 800 !important;
 }
 button {
-    border-radius: 16px !important;
-    font-weight: 900 !important;
-    box-shadow: 0 10px 24px rgba(15, 23, 42, 0.14) !important;
-}
-button.primary {
-    background: linear-gradient(90deg, #0f172a, #0f5f8f) !important;
+    background: linear-gradient(90deg, #020617, #1d4ed8) !important;
     color: #ffffff !important;
     border: none !important;
+    border-radius: 16px !important;
+    padding: 13px 24px !important;
+    font-weight: 900 !important;
+    font-size: 1rem !important;
+    box-shadow: 0 10px 22px rgba(37, 99, 235, 0.30) !important;
 }
-button.primary:hover {
-    background: linear-gradient(90deg, #020617, #075985) !important;
+button:hover {
+    background: linear-gradient(90deg, #000000, #2563eb) !important;
     transform: translateY(-1px);
 }
-#ocs-input textarea {
-    font-family: "Menaion", serif !important;
-    font-size: 24px !important;
-    line-height: 1.8 !important;
-    font-weight: normal !important;
+#output-box textarea {
+    background: #f8fbff !important;
+    border: 1.5px solid #1d4ed8 !important;
+    color: #020617 !important;
+    font-family: Consolas, "Courier New", monospace !important;
 }
-#sentence-box {
-    background: #f8fafc;
-    border: 1px solid #cbd5e1;
-    border-left: 6px solid #0f5f8f;
+#token-card {
+    background: #f8fbff;
+    border: 1.5px solid #1d4ed8;
     border-radius: 18px;
-    padding: 16px;
-    color: #0f172a;
-    font-size: 1.05rem;
-    line-height: 1.7;
-    margin-bottom: 18px;
+    padding: 18px;
+    box-shadow: 0 8px 20px rgba(37, 99, 235, 0.10);
 }
-#sentence-box b {
-    color: #0f172a;
-    font-family: Inter, Arial, sans-serif !important;
+.lemma-table {
+    width: 100%;
+    border-collapse: separate;
+    border-spacing: 0 8px;
     font-size: 1rem;
 }
-#sentence-text {
-    color: #0f172a;
-    font-family: "Menaion", serif !important;
-    font-size: 25px !important;
-    font-weight: normal !important;
-    line-height: 1.8 !important;
-    margin-top: 8px;
-}
-#table-title {
-    color: #0f5f8f;
+.lemma-table th {
+    background: linear-gradient(90deg, #020617, #1d4ed8);
+    color: white;
+    padding: 12px 14px;
+    text-align: left;
     font-weight: 900;
-    font-size: 1rem;
-    margin: 6px 0 10px 0;
 }
-#lemma-table,
-#lemma-table table,
-#lemma-table td,
-#lemma-table th,
-#lemma-table input,
-#lemma-table textarea,
-#lemma-table .cell-wrap,
-#lemma-table .table-wrap {
-    font-family: "Menaion", serif !important;
-    font-size: 21px !important;
-    font-weight: normal !important;
-    line-height: 1.6 !important;
+.lemma-table th:first-child {
+    border-radius: 12px 0 0 12px;
 }
-.dataframe {
-    border-radius: 16px !important;
-    overflow: hidden !important;
+.lemma-table th:last-child {
+    border-radius: 0 12px 12px 0;
+}
+.lemma-table td {
+    background: #ffffff;
+    color: #020617;
+    padding: 12px 14px;
+    border-top: 1px solid #bfdbfe;
+    border-bottom: 1px solid #bfdbfe;
+    font-weight: 700;
+}
+.lemma-table td:first-child {
+    border-left: 1px solid #bfdbfe;
+    border-radius: 12px 0 0 12px;
+}
+.lemma-table td:last-child {
+    border-right: 1px solid #bfdbfe;
+    border-radius: 0 12px 12px 0;
+    color: #1d4ed8;
+}
+#note {
+    color: #1e3a8a;
+    font-size: 0.92rem;
+    text-align: center;
+    margin-top: 1rem;
+    font-weight: 600;
 }
 footer,
 .api,
@@ -212,216 +199,464 @@ button[aria-label="Settings"] {
 }
 """
 
-HERO_HTML = """
-<div id="hero">
-    <h1>Combo Tool Demo</h1>
-    <p>
-        Paste Old Church Slavonic text. The app preprocesses the input, tokenizes it, and returns each token with its predicted lemma.
-    </p>
-</div>
-"""
 
+def make_html_table(tokens, lemmas):
+    if not tokens:
+        return ""
 
-@lru_cache(maxsize=1)
-def get_model_paths():
-    repo_dir = snapshot_download(
-        repo_id=MODEL_REPO_ID,
-        local_dir=HF_LOCAL_DIR,
-    )
+    rows = []
+
+    for token, lemma in zip(tokens, lemmas):
+        rows.append(
+            f"""
+            <tr>
+                <td>{html.escape(token)}</td>
+                <td>{html.escape(lemma)}</td>
+            </tr>
+            """
+        )
 
-    model_dir = Path(repo_dir)
+    return f"""
+    <div id="token-card">
+        <table class="lemma-table">
+            <thead>
+                <tr>
+                    <th>Word</th>
+                    <th>Lemma</th>
+                </tr>
+            </thead>
+            <tbody>
+                {''.join(rows)}
+            </tbody>
+        </table>
+    </div>
+    """
 
-    tokenizer_model = model_dir / "models" / MODEL_VARIANT / "tokenize" / "cu_proiel_tokenizer.pt"
-    lemma_model = model_dir / "models" / MODEL_VARIANT / "lemma" / "cu_proiel_nocharlm_lemmatizer.pt"
 
-    if not tokenizer_model.exists():
-        raise FileNotFoundError(f"Tokenizer model not found: {tokenizer_model}")
+def update_treebanks(language):
+    choices = treebank_choices(language)
 
-    if not lemma_model.exists():
-        raise FileNotFoundError(f"Lemmatizer model not found: {lemma_model}")
+    return gr.Dropdown(
+        choices=choices,
+        value=default_treebank(language),
+    )
 
-    return tokenizer_model, lemma_model
 
+def selected_model_id(display_name):
+    if not display_name or display_name not in DISPLAY_TO_ID:
+        raise ValueError("Please select a valid language and treebank.")
 
-@lru_cache(maxsize=1)
-def load_nlp_pipeline():
-    tokenizer_model, lemma_model = get_model_paths()
+    return DISPLAY_TO_ID[display_name]
 
-    stanza.download(
-        lang=LANG,
-        model_dir=STANZA_DIR,
-        processors={"pos": POS_PACKAGE},
-        package=None,
-        verbose=False,
-    )
 
-    return stanza.Pipeline(
-        lang=LANG,
-        dir=STANZA_DIR,
-        package=None,
-        processors={
-            "tokenize": TOKENIZER_PACKAGE,
-            "pos": POS_PACKAGE,
-            "lemma": LEMMATIZER_PACKAGE,
-        },
-        tokenize_model_path=str(tokenizer_model),
-        lemma_model_path=str(lemma_model),
-        tokenize_pretokenized=False,
-        use_gpu=True,
-        verbose=False,
+@lru_cache(maxsize=128)
+def load_vocab_chars_for_model(model_id):
+    
+    item = REGISTRY[model_id]
+
+    vocab_path = hf_hub_download(
+        repo_id=MODEL_REPO_ID,
+        repo_type="model",
+        filename=f"{MODEL_ROOT}/{item['folder']}/{item['vocab_file']}",
     )
 
+    with open(vocab_path, encoding="utf8") as f:
+        vocab_data = json.load(f)
 
-def clean_text(text):
-    return str(text or "").strip()
+    return set(vocab_data["char2idx"]) - {"<pad>", "<sos>", "<eos>", "<unk>"}
 
 
-def apply_mapping(text, mapping):
-    for old_char, new_char in sorted(mapping.items(), key=lambda item: len(item[0]), reverse=True):
-        text = text.replace(old_char, new_char)
-    return text
+def unsupported_input_for_model(text, allowed_chars, max_bad_ratio=0.60, min_checked_chars=4):
+    
+    checked = []
+    bad = []
 
+    for ch in text:
+        if ch.isspace():
+            continue
 
-def preprocess_text(text):
-    text = apply_mapping(text, REPLACEMENTS1)
-    return text
+        
+        if ch.isdigit() or ch in {".", ",", ";", ":", "!", "?", "-", "'", '"', "(", ")", "[", "]", "/"}:
+            continue
 
+        checked.append(ch)
 
-def make_sentence_html(text):
-    safe_text = (
-        text.replace("&", "&amp;")
-        .replace("<", "&lt;")
-        .replace(">", "&gt;")
-    )
+        if ch not in allowed_chars:
+            bad.append(ch)
 
-    return f"""
-    <div id="sentence-box">
-        <b>Input text</b>
-        <div id="sentence-text">{safe_text}</div>
-    </div>
-    """
+    if len(checked) < min_checked_chars:
+        return False, []
 
+    bad_ratio = len(bad) / len(checked)
 
-@spaces.GPU(duration=120)
-def analyze_text(text):
-    text = clean_text(text)
+    return bad_ratio >= max_bad_ratio, sorted(set(bad))
 
-    if not text:
-        empty_table = pd.DataFrame(columns=["Token", "Lemma"])
-        return (
-            gr.update(visible=True),
-            gr.update(value="<div id='sentence-box'>Please paste a sentence or raw text first.</div>"),
-            gr.update(value="", visible=False),
-            gr.update(value=empty_table, visible=False),
-        )
+
+def lemmatize_sentence(sentence, display_name):
+    sentence = str(sentence).strip()
+
+    if not sentence:
+        return "", ""
 
     try:
-        processed_text = preprocess_text(text)
-        nlp = load_nlp_pipeline()
-        doc = nlp(processed_text)
-
-        rows = []
-
-        for sentence in doc.sentences:
-            for word in sentence.words:
-                rows.append({
-                    "Token": word.text or "_",
-                    "Lemma": word.lemma or "_",
-                })
-
-        table = pd.DataFrame(rows, columns=["Token", "Lemma"])
-
-        return (
-            gr.update(visible=True),
-            gr.update(value=make_sentence_html(processed_text)),
-            gr.update(value="<div id='table-title'>Token and lemma output</div>", visible=True),
-            gr.update(value=table, visible=True),
+        model_id = selected_model_id(display_name)
+    except ValueError as e:
+        return "", str(e)
+
+    allowed_chars = load_vocab_chars_for_model(model_id)
+
+    is_bad, bad = unsupported_input_for_model(sentence, allowed_chars)
+    if is_bad:
+        return "",f"This input does not seem to match the selected language/treebank. Please select your desired language and treebank, then try again. Unsupported characters: {' '.join(bad[:20])}"
+
+    lemmatizer = load_lemmatizer(model_id, DEVICE)
+
+    tokens = sentence.split()
+    lemmas = lemmatizer.lemmatize_sentence(tokens)
+
+    lemmatized_sentence = " ".join(lemmas)
+    token_html = make_html_table(tokens, lemmas)
+
+    return lemmatized_sentence, token_html
+
+
+def parse_conllu_sentences_from_text(text):
+    text = text.strip()
+    sents = []
+
+    for block in re.split(r"\n\n+", text):
+        sent = []
+
+        for line in block.splitlines():
+            if not line or line.startswith("#"):
+                continue
+
+            cols = line.split("\t")
+
+            if len(cols) != 10:
+                continue
+
+            tok_id = cols[0]
+
+            if "-" in tok_id or "." in tok_id:
+                continue
+
+            form = cols[1]
+            sent.append(form)
+
+        if sent:
+            sents.append(sent)
+
+    return sents
+
+
+def make_source_for_token(tokens, index, k_context, sep_char):
+    form = tokens[index]
+
+    left_context = tokens[max(0, index - k_context):index]
+    right_context = tokens[index + 1:index + 1 + k_context]
+
+    left = " ".join(left_context).strip()
+    right = " ".join(right_context).strip()
+
+    src_left = left + " " if left else ""
+    src_right = " " + right if right else ""
+
+    return f"{src_left}{sep_char}{form}{sep_char}{src_right}"
+
+
+def make_all_sources_from_conllu(text, lemmatizer):
+    sents = parse_conllu_sentences_from_text(text)
+
+    sources = []
+
+    for tokens in sents:
+        for i in range(len(tokens)):
+            src_string = make_source_for_token(
+                tokens=tokens,
+                index=i,
+                k_context=lemmatizer.k_context,
+                sep_char=lemmatizer.sep_char,
+            )
+            sources.append(src_string)
+
+    return sources
+
+
+def predict_sources_batched(sources, lemmatizer, batch_size=BATCH_SIZE):
+    preds_all = []
+
+    if not sources:
+        return preds_all
+
+    pad_id = lemmatizer.vocab.char2idx["<pad>"]
+    sos_id = lemmatizer.vocab.char2idx["<sos>"]
+    eos_id = lemmatizer.vocab.char2idx["<eos>"]
+
+    for start in range(0, len(sources), batch_size):
+        batch_sources = sources[start:start + batch_size]
+
+        src_ids_list = []
+        src_lens = []
+
+        for src_string in batch_sources:
+            src_ids = (
+                [sos_id]
+                + lemmatizer.vocab.encode(src_string)
+                + [eos_id]
+            )
+
+            src_ids_list.append(src_ids)
+            src_lens.append(len(src_ids))
+
+        max_len = max(src_lens)
+
+        padded = [
+            ids + [pad_id] * (max_len - len(ids))
+            for ids in src_ids_list
+        ]
+
+        src = torch.tensor(
+            padded,
+            dtype=torch.long,
+            device=lemmatizer.device,
         )
 
-    except Exception as error:
-        empty_table = pd.DataFrame(columns=["Token", "Lemma"])
+        src_lens_tensor = torch.tensor(
+            src_lens,
+            dtype=torch.long,
+            device=lemmatizer.device,
+        )
 
-        return (
-            gr.update(visible=True),
-            gr.update(value=f"<div id='sentence-box'>Error: {str(error)}</div>"),
-            gr.update(value="", visible=False),
-            gr.update(value=empty_table, visible=False),
+        batch_preds = lemmatizer.model.generate(
+            src,
+            src_lens_tensor,
+            lemmatizer.vocab,
+            max_len=lemmatizer.max_gen_len,
         )
 
+        preds_all.extend(batch_preds)
+
+    return preds_all
+
 
-def clear_all():
-    empty_table = pd.DataFrame(columns=["Token", "Lemma"])
+def predict_conllu_lemmas(text, lemmatizer):
+    sources = make_all_sources_from_conllu(text, lemmatizer)
 
-    return (
-        "",
-        gr.update(visible=False),
-        gr.update(value=""),
-        gr.update(value="", visible=False),
-        gr.update(value=empty_table, visible=False),
+    return predict_sources_batched(
+        sources=sources,
+        lemmatizer=lemmatizer,
+        batch_size=BATCH_SIZE,
     )
 
 
-def load_example():
-    return EXAMPLE_TEXT
+def write_back_conllu(input_text, preds_all):
+    text = input_text.rstrip("\n")
+    blocks = re.split(r"\n\n+", text)
+
+    out_blocks = []
+    p = 0
+
+    for block in blocks:
+        lines = block.split("\n")
+        new_lines = []
+
+        for line in lines:
+            if not line or line.startswith("#"):
+                new_lines.append(line)
+                continue
+
+            cols = line.split("\t")
+
+            if len(cols) != 10:
+                new_lines.append(line)
+                continue
+
+            tok_id = cols[0]
+
+            if "-" in tok_id or "." in tok_id:
+                new_lines.append(line)
+                continue
+
+            pred = preds_all[p] if p < len(preds_all) else "_"
+            cols[TARGET_COL_IDX] = pred if pred else "_"
+
+            new_lines.append("\t".join(cols))
+            p += 1
+
+        out_blocks.append("\n".join(new_lines))
+
+    output_text = "\n\n".join(out_blocks).rstrip() + "\n\n"
+
+    return output_text, p
+
+
+def lemmatize_conllu_file(file_obj, display_name):
+    if file_obj is None:
+        return gr.update(value=None, visible=False), "Please upload a CoNLL-U file."
+
+    try:
+        model_id = selected_model_id(display_name)
+    except ValueError as e:
+        return gr.update(value=None, visible=False), str(e)
+
+    input_path = Path(file_obj.name)
+
+    with open(input_path, encoding="utf8") as f:
+        text = f.read()
+
+    lemmatizer = load_lemmatizer(model_id, DEVICE)
+
+    preds = predict_conllu_lemmas(text, lemmatizer)
+    output_text, total = write_back_conllu(text, preds)
+
+    safe_model_name = display_name.replace(" ", "_").replace("-", "_")
+    safe_model_name = re.sub(r"[^A-Za-z0-9_]+", "", safe_model_name)
+
+    out_path = (
+        Path(tempfile.gettempdir())
+        / f"{input_path.stem}.{safe_model_name}.lemmatized.conllu"
+    )
+
+    with open(out_path, "w", encoding="utf8", newline="\n") as f:
+        f.write(output_text)
+
+    message = (
+        f"Done. Wrote {total:,} lemma predictions.\n"
+        f"Input used: FORM column only.\n"
+        f"Updated column: LEMMA only.\n"
+        f"All other CoNLL-U columns and comments were preserved."
+    )
+
+    return gr.update(value=str(out_path), visible=True), message
+
+
+def reset_download_button(file_obj):
+    return gr.update(value=None, visible=False), ""
+
+
+DEFAULT_LANGUAGE = default_language()
+DEFAULT_TREEBANK = default_treebank(DEFAULT_LANGUAGE)
 
 
 with gr.Blocks(
-    title="Combo New Tool",
+    title="oldslaviclemma",
     css=CUSTOM_CSS,
     theme=gr.themes.Soft(
-        primary_hue="cyan",
-        secondary_hue="blue",
+        primary_hue="blue",
+        secondary_hue="sky",
         neutral_hue="slate",
     ),
 ) as demo:
-    gr.HTML(HERO_HTML)
-
     with gr.Column(elem_id="main-card"):
-        text_input = gr.Textbox(
-            label="Input text",
-            lines=8,
-            placeholder="Paste Old Church Slavonic text here...",
-            elem_id="ocs-input",
+        gr.Markdown("# oldslaviclemma", elem_id="title")
+
+        gr.Markdown(
+            "Select a language and treebank. Paste one sentence or upload a tokenized CoNLL-U file. "
+            "The system returns lemma predictions while preserving the original tokenization.",
+            elem_id="subtitle",
+        )
+
+        gr.HTML(
+            """
+            <div id="badge-row">
+                <span>oldslaviclemma</span>
+                <span>60+ languages</span>
+                <span>110+ treebanks</span>
+                <span>UD v2.12</span>
+                <span>Lemmatization</span>
+            </div>
+            """
         )
 
         with gr.Row():
-            analyze_button = gr.Button("Lemmatize", variant="primary")
-            example_button = gr.Button("Load example")
-            clear_button = gr.Button("Clear")
-
-    with gr.Column(elem_id="output-card", visible=False) as output_panel:
-        sentence_output = gr.HTML()
-        table_title = gr.HTML(visible=False)
-
-        lemma_table = gr.Dataframe(
-            headers=["Token", "Lemma"],
-            interactive=False,
-            wrap=True,
-            show_label=False,
-            visible=False,
-            elem_id="lemma-table",
+            language_input = gr.Dropdown(
+                label="Language",
+                choices=LANGUAGES,
+                value=DEFAULT_LANGUAGE,
+            )
+
+            treebank_input = gr.Dropdown(
+                label="Treebank",
+                choices=treebank_choices(DEFAULT_LANGUAGE),
+                value=DEFAULT_TREEBANK,
+            )
+
+        with gr.Tab("Sentence input"):
+            sentence_input = gr.Textbox(
+                label="Input sentence",
+                lines=5,
+                value="",
+                placeholder="Paste a sentence with words separated by spaces...",
+            )
+
+            run_button = gr.Button("Lemmatize sentence")
+
+            sentence_output = gr.Textbox(
+                label="Lemmatized sentence",
+                lines=5,
+                elem_id="output-box",
+            )
+
+            token_output = gr.HTML(
+                label="Word-level output",
+                value="",
+            )
+
+            gr.Markdown(
+                "Please paste one sentence with whitespace tokenization.",
+                elem_id="note",
+            )
+
+        with gr.Tab("CoNLL-U file input"):
+            gr.Markdown(
+                "Upload an already-tokenized CoNLL-U file. "
+                "The app reads the FORM column, predicts the LEMMA column, "
+                "and preserves comments, token IDs, UPOS, XPOS, FEATS, HEAD, DEPREL, DEPS, and MISC."
+            )
+
+            conllu_input = gr.File(
+                label="Upload CoNLL-U file",
+                file_types=[".conllu", ".txt"],
+            )
+
+            conllu_button = gr.Button("Lemmatize CoNLL-U file")
+
+            conllu_output = gr.DownloadButton(
+                label="Download lemmatized CoNLL-U file",
+                value=None,
+                visible=False,
+            )
+
+            conllu_message = gr.Textbox(
+                label="Status",
+                lines=4,
+            )
+
+        language_input.change(
+            fn=update_treebanks,
+            inputs=language_input,
+            outputs=treebank_input,
         )
 
-    analyze_button.click(
-        fn=analyze_text,
-        inputs=text_input,
-        outputs=[output_panel, sentence_output, table_title, lemma_table],
-        show_progress="hidden",
-    )
+        run_button.click(
+            fn=lemmatize_sentence,
+            inputs=[sentence_input, treebank_input],
+            outputs=[sentence_output, token_output],
+        )
 
-    example_button.click(
-        fn=load_example,
-        inputs=None,
-        outputs=text_input,
-        show_progress="hidden",
-    )
+        conllu_input.change(
+            fn=reset_download_button,
+            inputs=conllu_input,
+            outputs=[conllu_output, conllu_message],
+        )
 
-    clear_button.click(
-        fn=clear_all,
-        inputs=None,
-        outputs=[text_input, output_panel, sentence_output, table_title, lemma_table],
-        show_progress="hidden",
-    )
+        conllu_button.click(
+            fn=lemmatize_conllu_file,
+            inputs=[conllu_input, treebank_input],
+            outputs=[conllu_output, conllu_message],
+        )
 
 
-demo.launch()
+demo.queue()
+demo.launch(server_name="0.0.0.0", server_port=7860, ssr_mode=False)

inference.py

@@ -0,0 +1,150 @@
+import json
+from functools import lru_cache
+from pathlib import Path
+
+import torch
+from huggingface_hub import hf_hub_download
+
+try:
+    from model import LemmaModel, Vocab
+except ImportError:
+    from .model import LemmaModel, Vocab
+
+
+MODEL_REPO_ID = "usmannawaz/oldslaviclemma"
+MODEL_ROOT = "oldslaviclemma212"
+
+
+@lru_cache(maxsize=1)
+def load_registry(registry_path="models_registry.json"):
+    
+    registry_path = Path(registry_path)
+    with registry_path.open(encoding="utf8") as f:
+        return json.load(f)
+
+
+class OldSlavicLemmatizer:
+    def __init__(self, model, vocab, config, device):
+        self.model = model
+        self.vocab = vocab
+        self.config = config
+        self.device = torch.device(device)
+        self.sep_char = config.get("sep_char", "⟂")
+        self.k_context = int(config.get("k_context", 2))
+        self.max_gen_len = int(config.get("max_gen_len", 30))
+
+    def make_source(self, form, left_context=None, right_context=None):
+        left_context = left_context or []
+        right_context = right_context or []
+
+        left = " ".join(left_context[-self.k_context:]).strip()
+        right = " ".join(right_context[:self.k_context]).strip()
+
+        src_left = left + " " if left else ""
+        src_right = " " + right if right else ""
+
+        return f"{src_left}{self.sep_char}{form}{self.sep_char}{src_right}"
+
+    def lemmatize(self, form, left_context=None, right_context=None):
+        src_string = self.make_source(
+            form=form,
+            left_context=left_context,
+            right_context=right_context,
+        )
+
+        src_ids = (
+            [self.vocab.char2idx["<sos>"]]
+            + self.vocab.encode(src_string)
+            + [self.vocab.char2idx["<eos>"]]
+        )
+
+        src = torch.tensor([src_ids], dtype=torch.long, device=self.device)
+        src_lens = torch.tensor([len(src_ids)], dtype=torch.long, device=self.device)
+
+        return self.model.generate(
+            src,
+            src_lens,
+            self.vocab,
+            max_len=self.max_gen_len,
+        )[0]
+
+    def lemmatize_sentence(self, tokens):
+        lemmas = []
+
+        for i, token in enumerate(tokens):
+            left_context = tokens[max(0, i - self.k_context):i]
+            right_context = tokens[i + 1:i + 1 + self.k_context]
+
+            lemma = self.lemmatize(
+                token,
+                left_context=left_context,
+                right_context=right_context,
+            )
+
+            lemmas.append(lemma)
+
+        return lemmas
+
+
+@lru_cache(maxsize=3)
+def load_lemmatizer(model_id, device=None):
+    
+    if device is None:
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+    device = torch.device(device)
+
+    registry = load_registry("models_registry.json")
+    if model_id not in registry:
+        raise KeyError(f"Model id not found in registry: {model_id}")
+
+    item = registry[model_id]
+    folder = item["folder"]
+
+    config_path = hf_hub_download(
+        repo_id=MODEL_REPO_ID,
+        repo_type="model",
+        filename=f"{MODEL_ROOT}/{folder}/{item['config_file']}",
+    )
+
+    vocab_path = hf_hub_download(
+        repo_id=MODEL_REPO_ID,
+        repo_type="model",
+        filename=f"{MODEL_ROOT}/{folder}/{item['vocab_file']}",
+    )
+
+    weights_path = hf_hub_download(
+        repo_id=MODEL_REPO_ID,
+        repo_type="model",
+        filename=f"{MODEL_ROOT}/{folder}/{item['model_file']}",
+    )
+
+    with open(config_path, encoding="utf8") as f:
+        config = json.load(f)
+
+    with open(vocab_path, encoding="utf8") as f:
+        vocab_data = json.load(f)
+
+    vocab = Vocab(
+        char2idx=vocab_data["char2idx"],
+        idx2char=vocab_data["idx2char"],
+    )
+
+    model = LemmaModel(
+        vocab_size=len(vocab.char2idx),
+        char_emb_dim=int(config["char_emb_dim"]),
+        hidden_size=int(config["hidden_size"]),
+        drop_prob=float(config["drop_prob"]),
+        num_heads=int(config["num_heads"]),
+        max_gen_len=int(config.get("max_gen_len", 30)),
+    ).to(device)
+
+    state = torch.load(weights_path, map_location=device)
+    model.load_state_dict(state)
+    model.eval()
+
+    return OldSlavicLemmatizer(
+        model=model,
+        vocab=vocab,
+        config=config,
+        device=device,
+    )

model.py

@@ -0,0 +1,133 @@
+from typing import List
+import torch
+import torch.nn as nn
+
+
+class Vocab:
+    def __init__(self, char2idx=None, idx2char=None):
+        if char2idx is None:
+            char2idx = {
+                "<pad>": 0,
+                "<sos>": 1,
+                "<eos>": 2,
+                "<unk>": 3,
+            }
+        self.char2idx = char2idx
+        if idx2char is None:
+            self.idx2char = {i: c for c, i in self.char2idx.items()}
+        else:
+            self.idx2char = {int(k): v for k, v in idx2char.items()}
+
+    def encode(self, s: str) -> List[int]:
+        unk = self.char2idx["<unk>"]
+        return [self.char2idx.get(ch, unk) for ch in s]
+
+    def decode(self, ids: List[int]) -> str:
+        out = []
+        eos_id = self.char2idx["<eos>"]
+        for i in ids:
+            if i == eos_id:
+                break
+            if i > eos_id:
+                out.append(self.idx2char.get(int(i), ""))
+        return "".join(out)
+
+
+class LemmaModel(nn.Module):
+    def __init__(
+        self,
+        vocab_size: int,
+        char_emb_dim: int = 96,
+        hidden_size: int = 128,
+        drop_prob: float = 0.30,
+        num_heads: int = 16,
+        max_gen_len: int = 30,
+    ):
+        super().__init__()
+        self.max_gen_len = max_gen_len
+        self.emb = nn.Embedding(vocab_size, char_emb_dim, padding_idx=0)
+        self.dropout_enc = nn.Dropout(drop_prob)
+        self.dropout_dec = nn.Dropout(drop_prob)
+        self.dropout_att = nn.Dropout(drop_prob)
+        self.enc1 = nn.LSTM(char_emb_dim, hidden_size, bidirectional=True, batch_first=True)
+        self.enc2 = nn.LSTM(hidden_size * 2, hidden_size, bidirectional=True, batch_first=True)
+        self.attn = nn.MultiheadAttention(hidden_size * 2, num_heads, batch_first=True)
+        self.dec = nn.LSTM(char_emb_dim + hidden_size * 4, hidden_size * 2, batch_first=True)
+        self.dec_cross_attn = nn.MultiheadAttention(
+            embed_dim=hidden_size * 2,
+            num_heads=num_heads,
+            kdim=hidden_size * 4,
+            vdim=hidden_size * 4,
+            batch_first=True,
+        )
+        self.out = nn.Linear(hidden_size * 2, vocab_size, bias=True)
+
+    def encode(self, src, src_lens):
+        emb = self.emb(src)
+        packed1 = nn.utils.rnn.pack_padded_sequence(
+            emb, src_lens.cpu(), batch_first=True, enforce_sorted=False
+        )
+        enc1_o, _ = self.enc1(packed1)
+        enc1_o, _ = nn.utils.rnn.pad_packed_sequence(enc1_o, batch_first=True)
+        enc1_o = self.dropout_enc(enc1_o)
+        packed2 = nn.utils.rnn.pack_padded_sequence(
+            enc1_o, src_lens.cpu(), batch_first=True, enforce_sorted=False
+        )
+        enc2_o, _ = self.enc2(packed2)
+        enc2_o, _ = nn.utils.rnn.pad_packed_sequence(enc2_o, batch_first=True)
+        enc2_o = self.dropout_enc(enc2_o)
+        attn_o, _ = self.attn(enc1_o, enc2_o, enc2_o)
+        attn_o = self.dropout_att(attn_o)
+        return torch.cat([enc2_o, attn_o], dim=-1)
+
+    def forward(self, src, src_lens, tgt):
+        encoder_combined = self.encode(src, src_lens)
+        dt = self.emb(tgt[:, :-1])
+        target_len = dt.size(1)
+        if encoder_combined.size(1) >= target_len:
+            comb_trim = encoder_combined[:, :target_len, :]
+        else:
+            pad = encoder_combined.new_zeros(
+                encoder_combined.size(0),
+                target_len - encoder_combined.size(1),
+                encoder_combined.size(2),
+            )
+            comb_trim = torch.cat([encoder_combined, pad], dim=1)
+        dec_inp = torch.cat([dt, comb_trim], dim=-1)
+        dec_o, _ = self.dec(dec_inp)
+        dec_o = self.dropout_dec(dec_o)
+        cross_out, _ = self.dec_cross_attn(dec_o, encoder_combined, encoder_combined)
+        cross_out = self.dropout_att(cross_out)
+        return self.out(cross_out)
+
+    def generate(self, src, src_lens, vocab, max_len=None):
+        self.eval()
+        if max_len is None:
+            max_len = self.max_gen_len
+        batch_size = src.size(0)
+        with torch.no_grad():
+            encoder_combined = self.encode(src, src_lens)
+            source_len = encoder_combined.size(1)
+            cur = torch.full(
+                (batch_size, 1), vocab.char2idx["<sos>"], device=src.device, dtype=torch.long
+            )
+            hidden = None
+            hyps = [[] for _ in range(batch_size)]
+            for step in range(max_len):
+                emb_t = self.emb(cur).squeeze(1)
+                if source_len == 0:
+                    comb_t = encoder_combined[:, 0, :]
+                else:
+                    comb_t = encoder_combined[:, min(step, source_len - 1), :]
+                dec_inp_t = torch.cat([emb_t, comb_t], dim=-1).unsqueeze(1)
+                dec_o, hidden = self.dec(dec_inp_t, hidden)
+                dec_o = self.dropout_dec(dec_o)
+                cross_out, _ = self.dec_cross_attn(dec_o, encoder_combined, encoder_combined)
+                cross_out = self.dropout_att(cross_out)
+                logits = self.out(cross_out)
+                cur = logits.argmax(-1)
+                for i in range(batch_size):
+                    hyps[i].append(int(cur[i, 0].item()))
+                if all(int(cur[i, 0].item()) == vocab.char2idx["<eos>"] for i in range(batch_size)):
+                    break
+        return [vocab.decode(h) for h in hyps]

requirements.txt

@@ -1,5 +1,3 @@
+torch
 gradio
-stanza
-huggingface_hub
-pandas
-torch
+huggingface_hub