Update app.py

app.py

@@ -1,57 +1,41 @@
 import os
 import re
-import importlib.util
 from pathlib import Path
 
 import torch
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
+from peft import PeftModel  # NEW
 
-MODEL_ID = "ZennyKenny/oss-20b-prereform-to-modern-ru-merged"
+MODEL_ID_BASE = "openai/gpt-oss-20b"  # base model
+ADAPTER_REPO = "ZennyKenny/oss-20b-prereform-to-modern-ru-merged"
+ADAPTER_SUBFOLDER = "checkpoint-60"   # where adapter lives in your repo
 
-# ----------------- Load SYSTEM_PROMPT from 'text-prompt.py' -----------------
+# ---- load SYSTEM_PROMPT from text-prompt.py (same as before) ----
 def _load_system_prompt():
     prompt_path = Path(__file__).with_name("text-prompt.py")
-    default_prompt = (
-        "Ты компетентный редактор русского языка. "
-        "Преобразуй дореформенную русскую орфографию (до 1918 года) "
-        "в современную орфографию. Сохраняй смысл, пунктуацию и регистр. "
-        "Не добавляй комментариев. Верни только преобразованный текст."
-    )
+    default = ("Ты компетентный редактор русского языка. "
+               "Преобразуй дореформенную русскую орфографию (до 1918 года) "
+               "в современную орфографию. Сохраняй смысл, пунктуацию и регистр. "
+               "Не добавляй комментариев. Верни только преобразованный текст.")
     try:
-        if not prompt_path.exists():
-            return default_prompt
-        spec = importlib.util.spec_from_file_location("text_prompt_mod", str(prompt_path))
-        mod = importlib.util.module_from_spec(spec)
-        assert spec and spec.loader, "Cannot load spec for text-prompt.py"
-        spec.loader.exec_module(mod)  # type: ignore[attr-defined]
-        return getattr(mod, "SYSTEM_PROMPT", default_prompt)
+        ns = {}
+        exec(prompt_path.read_text(encoding="utf-8"), ns) if prompt_path.exists() else None
+        return ns.get("SYSTEM_PROMPT", default)
     except Exception:
-        return default_prompt
+        return default
 
 SYSTEM_PROMPT = _load_system_prompt()
 
-# ----------------- Fallback: rule-based converter (no ML needed) -----------------
-REPLACEMENTS = [
-    ("Ѣ", "Е"), ("ѣ", "е"),
-    ("І", "И"), ("і", "и"),
-    ("Ѳ", "Ф"), ("ѳ", "ф"),
-    ("Ѵ", "И"), ("ѵ", "и"),
-]
-TERMINAL_HARD_SIGN = re.compile(r"(?i)ъ\b")  # remove word-final hard sign
-MULTI_SPACES = re.compile(r"[ \t]{2,}")
-
-def rule_based_convert(text: str) -> str:
-    if not text:
-        return ""
-    out = text
-    for old, new in REPLACEMENTS:
-        out = out.replace(old, new)
-    out = TERMINAL_HARD_SIGN.sub("", out)
-    out = MULTI_SPACES.sub(" ", out)
-    return out
+# ---- simple rule-based fallback (unchanged) ----
+REPLACEMENTS = [("Ѣ","Е"),("ѣ","е"),("І","И"),("і","и"),("Ѳ","Ф"),("ѳ","ф"),("Ѵ","И"),("ѵ","и")]
+TERMINAL_HARD_SIGN = re.compile(r"(?i)ъ\b")
+def rule_based_convert(t):
+    if not t: return ""
+    for a,b in REPLACEMENTS: t = t.replace(a,b)
+    return TERMINAL_HARD_SIGN.sub("", t)
 
-# ----------------- Model state (CPU-only) -----------------
+# ---- model state (CPU only) ----
 _tokenizer = None
 _model = None
 _streamer = None
@@ -59,54 +43,42 @@ _MODEL_READY = False
 _MODEL_ERROR = None
 
 def build_prompt(text: str) -> str:
-    return (
-        f"{SYSTEM_PROMPT}\n\n"
-        f"Текст (дореформ.):\n{text.strip()}\n\n"
-        f"Текст (современная орфография):"
-    )
+    return f"{SYSTEM_PROMPT}\n\nТекст (дореформ.):\n{text.strip()}\n\nТекст (современная орфография):"
 
 def load_model_cpu():
-    """Force CPU load. Gracefully degrade if loading fails."""
+    """Load base model, then apply LoRA adapter from your repo."""
     global _tokenizer, _model, _streamer, _MODEL_READY, _MODEL_ERROR
     if _MODEL_READY or _MODEL_ERROR:
         return
-
     if os.getenv("DISABLE_MODEL", "0") == "1":
         _MODEL_ERROR = "Model disabled via DISABLE_MODEL=1."
         return
-
     try:
         os.environ["CUDA_VISIBLE_DEVICES"] = ""
-        os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
-
-        _tokenizer = AutoTokenizer.from_pretrained(
-            MODEL_ID, use_fast=True, trust_remote_code=True
-        )
-        _model = AutoModelForCausalLM.from_pretrained(
-            MODEL_ID,
+        _tokenizer = AutoTokenizer.from_pretrained(ADAPTER_REPO, use_fast=True, trust_remote_code=True)
+        base = AutoModelForCausalLM.from_pretrained(
+            MODEL_ID_BASE,
             trust_remote_code=True,
-            torch_dtype=torch.float32,      # CPU dtype
+            torch_dtype=torch.float32,
             low_cpu_mem_usage=True,
-            device_map=None,                # ensure CPU
+            device_map=None,
         ).to("cpu")
+        # Apply LoRA adapter from your repo/subfolder
+        _model = PeftModel.from_pretrained(base, ADAPTER_REPO, subfolder=ADAPTER_SUBFOLDER)
+        # (Optional) Merge for faster inference on CPU:
+        try:
+            _model = _model.merge_and_unload()
+        except Exception:
+            pass
         _streamer = TextStreamer(_tokenizer, skip_prompt=True, skip_special_tokens=True)
         _MODEL_READY = True
     except Exception as e:
         _MODEL_ERROR = f"{type(e).__name__}: {e}"
 
-def convert_with_model(
-    text: str,
-    max_new_tokens: int,
-    temperature: float,
-    top_p: float,
-    top_k: int,
-    repetition_penalty: float,
-    do_stream: bool
-) -> str:
+def convert_with_model(text, max_new_tokens, temperature, top_p, top_k, repetition_penalty, do_stream):
     prompt = build_prompt(text)
     inputs = _tokenizer(prompt, return_tensors="pt")
     input_ids = inputs.input_ids.to("cpu")
-
     gen_kwargs = dict(
         max_new_tokens=int(max_new_tokens),
         temperature=float(temperature),
@@ -115,65 +87,47 @@ def convert_with_model(
         repetition_penalty=float(repetition_penalty),
         do_sample=True,
     )
-
     if do_stream:
         chunks = []
-
-        class _BufStreamer(TextStreamer):
-            def on_finalized_text(self, text, stream_end=False):
-                chunks.append(text)
-
-        buf_streamer = _BufStreamer(_tokenizer, skip_prompt=True, skip_special_tokens=True)
-        _ = _model.generate(input_ids=input_ids, streamer=buf_streamer, **gen_kwargs)
+        class _Buf(TextStreamer):
+            def on_finalized_text(self, txt, stream_end=False):
+                chunks.append(txt)
+        buf = _Buf(_tokenizer, skip_prompt=True, skip_special_tokens=True)
+        _ = _model.generate(input_ids=input_ids, streamer=buf, **gen_kwargs)
         out = "".join(chunks)
     else:
         with torch.no_grad():
-            output_ids = _model.generate(input_ids=input_ids, **gen_kwargs)
-        out = _tokenizer.decode(output_ids[0], skip_special_tokens=True)
-
+            out_ids = _model.generate(input_ids=input_ids, **gen_kwargs)
+        out = _tokenizer.decode(out_ids[0], skip_special_tokens=True)
     marker = "Текст (современная орфография):"
     return out.split(marker, 1)[-1].strip() if marker in out else out.strip()
 
 def convert(text, max_new_tokens, temperature, top_p, top_k, repetition_penalty, do_stream):
     if not text or not text.strip():
         return ""
-
     load_model_cpu()
-
     if _MODEL_READY:
         try:
-            return convert_with_model(
-                text, max_new_tokens, temperature, top_p, top_k, repetition_penalty, do_stream
-            )
+            return convert_with_model(text, max_new_tokens, temperature, top_p, top_k, repetition_penalty, do_stream)
         except Exception:
             return rule_based_convert(text) + "\n\n[Примечание: использовано правило-базовое преобразование из-за ошибки генерации на CPU.]"
-    else:
-        note = "\n\n[Примечание: используется правило-базовое преобразование"
-        if _MODEL_ERROR:
-            note += f" (модель недоступна: {_MODEL_ERROR})"
-        note += ".]"
-        return rule_based_convert(text) + note
+    note = "\n\n[Примечание: используется правило-базовое преобразование"
+    if _MODEL_ERROR: note += f" (модель недоступна: {_MODEL_ERROR})"
+    note += ".]"
+    return rule_based_convert(text) + note
 
-# ----------------- UI -----------------
+# ---- Gradio UI (same structure as before) ----
 with gr.Blocks(title="Pre-reform → Modern Russian (CPU-only)") as demo:
     gr.Markdown(
         """
         # Преобразование дореформенной орфографии → современная (CPU-only)
-        Вставьте дореформенный русский текст — получите современную орфографию.  
-        Модель: `ZennyKenny/oss-20b-prereform-to-modern-ru-merged`
-
-        *Подсказка:* На CPU загрузка большой модели может быть недоступна; в таком случае
-        автоматически используется быстрый правило-базовый конвертер (ѣ→е, і→и, ѳ→ф, ѵ→и, удаление конечного ъ).
+        Модель: LoRA-адаптер к `openai/gpt-oss-20b` из `ZennyKenny/oss-20b-prereform-to-modern-ru-merged`.
+        При недоступности модели используется правило-базовый конвертер (ѣ→е, і→и, ѳ→ф, ѵ→и, удаление конечного ъ).
         """
     )
-
     with gr.Row():
-        with gr.Column(scale=1):
-            inp = gr.Textbox(
-                label="Ввод: дореформенный текст",
-                placeholder="Например: \"въ мирѣ сёмъ многа есть...\"",
-                lines=10
-            )
+        with gr.Column():
+            inp = gr.Textbox(label="Ввод: дореформенный текст", lines=10)
             with gr.Accordion("Параметры генерации (медленно на CPU)", open=False):
                 max_new_tokens = gr.Slider(8, 256, value=128, step=8, label="max_new_tokens")
                 temperature = gr.Slider(0.0, 1.2, value=0.2, step=0.05, label="temperature")
@@ -181,26 +135,17 @@ with gr.Blocks(title="Pre-reform → Modern Russian (CPU-only)") as demo:
                 top_k = gr.Slider(0, 100, value=40, step=1, label="top_k")
                 repetition_penalty = gr.Slider(1.0, 2.0, value=1.05, step=0.01, label="repetition_penalty")
                 do_stream = gr.Checkbox(value=False, label="Стриминг вывода")
-
             btn = gr.Button("Преобразовать", variant="primary")
-
-        with gr.Column(scale=1):
+        with gr.Column():
             out = gr.Textbox(label="Вывод: современная орфография", lines=12)
-
-    examples = [
-        ["въ семъ домѣ обитало три семейства, и каждое имѣло свои обыкновенія."],
-        ["Онъ шёлъ по узкой улѣцѣ, разсматривая вывѣски лавокъ и фонари."],
-        ["въ мирѣ сёмъ многа есть, чего мудрецу и не снилось."]
-    ]
-    gr.Examples(examples=examples, inputs=[inp])
-
-    def _on_click(text, a, b, c, d, e, f):
-        return convert(text, a, b, c, d, e, f)
-
+    gr.Examples(
+        examples=[["въ семъ домѣ обитало три семейства, и каждое имѣло свои обыкновенія."]],
+        inputs=[inp],
+    )
     btn.click(
-        _on_click,
+        lambda t,a,b,c,d,e,f: convert(t,a,b,c,d,e,f),
         inputs=[inp, max_new_tokens, temperature, top_p, top_k, repetition_penalty, do_stream],
-        outputs=[out]
+        outputs=[out],
     )
 
 if __name__ == "__main__":