Update app.py

app.py

@@ -1,24 +1,20 @@
 import os
-import re
 from pathlib import Path
 
 import gradio as gr
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
+from transformers import AutoTokenizer, AutoModelForCausalLM
 from peft import PeftModel
 import spaces  # ZeroGPU
 
 
 # ========= Config =========
-# Your LoRA repo and base model:
-MODEL_ID_BASE = "openai/gpt-oss-20b"  # base architecture
-ADAPTER_REPO = "ZennyKenny/oss-20b-prereform-to-modern-ru-merged"
-ADAPTER_SUBFOLDER = "checkpoint-60"   # LoRA lives here in your repo
+# Base model + your LoRA adapter (override via Space Secrets if needed)
+MODEL_ID_BASE = os.getenv("BASE_MODEL_ID", "openai/gpt-oss-20b")
+ADAPTER_REPO = os.getenv("ADAPTER_REPO", "ZennyKenny/oss-20b-prereform-to-modern-ru-merged")
+ADAPTER_SUBFOLDER = os.getenv("ADAPTER_SUBFOLDER", "checkpoint-60")  # change if your adapter folder differs
 
-# ZeroGPU toggle (you can also set in Space Secrets):
-USE_ZEROGPU = os.getenv("USE_ZEROGPU", "1") == "1"
-
-# ========= Load external prompt =========
+# ========= Load external system prompt =========
 def _load_system_prompt():
     path = Path(__file__).with_name("text-prompt.py")
     default = (
@@ -30,7 +26,7 @@ def _load_system_prompt():
     try:
         ns = {}
         if path.exists():
-            exec(path.read_text(encoding="utf-8"), ns)
+            exec(path.read_text(encoding='utf-8'), ns)
         return ns.get("SYSTEM_PROMPT", default)
     except Exception:
         return default
@@ -44,34 +40,17 @@ def build_prompt(text: str) -> str:
         f"Текст (современная орфография):"
     )
 
-# ========= Rule-based CPU fallback =========
-REPLACEMENTS = [
-    ("Ѣ", "Е"), ("ѣ", "е"),
-    ("І", "И"), ("і", "и"),
-    ("Ѳ", "Ф"), ("ѳ", "ф"),
-    ("Ѵ", "И"), ("ѵ", "и"),
-]
-TERMINAL_HARD_SIGN = re.compile(r"(?i)ъ\b")
-
-def rule_based_convert(text: str) -> str:
-    if not text:
-        return ""
-    for old, new in REPLACEMENTS:
-        text = text.replace(old, new)
-    text = TERMINAL_HARD_SIGN.sub("", text)
-    return text
-
-
-# ========= ZeroGPU path (model loads INSIDE the GPU-decorated function) =========
-# Note: Gradio/Spaces allocate the GPU ONLY during the call to this function.
-# Keep everything self-contained here: tokenizer, model, generate, return.
-@spaces.GPU(duration=180)  # allocate GPU just for this call (extend duration if you expect long runs)
+# ========= ZeroGPU inference =========
+@spaces.GPU(duration=180)  # GPU is leased only while this function runs
 def _infer_zerogpu(prompt: str, gen_kwargs: dict) -> str:
-    # Load tokenizer from your adapter repo (it contains tokenizer files)
+    # Tokenizer from adapter repo (it contains tokenizer files)
     tokenizer = AutoTokenizer.from_pretrained(ADAPTER_REPO, use_fast=True, trust_remote_code=True)
 
-    # Load base model on GPU (ZeroGPU provides an H200/A100-like device)
-    # Use bf16 if available, fallback fp16.
+    # Ensure pad token exists; if not, align it with EOS (common for GPT-like)
+    if tokenizer.pad_token_id is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    # Load base model on GPU with appropriate dtype
     torch_dtype = torch.bfloat16 if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else torch.float16
     base = AutoModelForCausalLM.from_pretrained(
         MODEL_ID_BASE,
@@ -83,26 +62,49 @@ def _infer_zerogpu(prompt: str, gen_kwargs: dict) -> str:
     # Apply LoRA adapter from your repo/subfolder
     model = PeftModel.from_pretrained(base, ADAPTER_REPO, subfolder=ADAPTER_SUBFOLDER)
 
-    # (Optional) Merge LoRA for faster generation and less VRAM fragmentation
+    # Optional: merge LoRA for faster generation
     try:
         model = model.merge_and_unload()
     except Exception:
         pass
 
-    # Generate on GPU
-    input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(model.device)
+    # Sync pad_token_id to model config to avoid warnings
+    try:
+        model.config.pad_token_id = tokenizer.pad_token_id
+    except Exception:
+        pass
+
+    # ----- Tokenize & always pass attention_mask -----
+    enc = tokenizer(prompt, return_tensors="pt", add_special_tokens=False)
+    input_ids = enc["input_ids"].to(model.device)
+    attention_mask = enc.get("attention_mask", torch.ones_like(input_ids)).to(model.device)
+
+    # Reasonable defaults
+    gen_kwargs = dict(gen_kwargs or {})
+    gen_kwargs.setdefault("use_cache", True)
+
+    # ----- Generate -----
     with torch.no_grad():
-        if "streamer" in gen_kwargs:
-            gen_kwargs.pop("streamer", None)
-        out_ids = model.generate(input_ids=input_ids, **gen_kwargs)
-    out = tokenizer.decode(out_ids[0], skip_special_tokens=True)
+        out_ids = model.generate(
+            input_ids=input_ids,
+            attention_mask=attention_mask,  # Key fix for pad==eos
+            **gen_kwargs,
+        )
 
-    marker = "Текст (современная орфография):"
-    return out.split(marker, 1)[-1].strip() if marker in out else out.strip()
+    # Decode ONLY the continuation (exclude prompt tokens)
+    continuation = out_ids[0, input_ids.shape[1]:]
+    out = tokenizer.decode(continuation, skip_special_tokens=True).strip()
 
+    # Fallback to full decode if continuation is empty (still no letter-replacement fallback)
+    if not out:
+        full = tokenizer.decode(out_ids[0], skip_special_tokens=True).strip()
+        marker = "Текст (современная орфография):"
+        out = full.split(marker, 1)[-1].strip() if marker in full else full
+
+    return out
 
 # ========= Orchestrator =========
-def convert(text, max_new_tokens, temperature, top_p, top_k, repetition_penalty, do_stream):
+def convert(text, max_new_tokens, temperature, top_p, top_k, repetition_penalty):
     if not text or not text.strip():
         return ""
 
@@ -116,25 +118,18 @@ def convert(text, max_new_tokens, temperature, top_p, top_k, repetition_penalty,
         do_sample=True,
     )
 
-    # Prefer ZeroGPU if enabled; otherwise CPU fallback
-    if USE_ZEROGPU:
-        try:
-            return _infer_zerogpu(prompt, gen_kwargs)
-        except Exception as e:
-            # If ZeroGPU is unavailable/rate limited/errored, gracefully fall back.
-            return rule_based_convert(text) + f"\n\n[Примечание: ZeroGPU недоступен или ошибка: {type(e).__name__}: {e}]"
-    else:
-        # Explicit CPU-only mode (fast fallback)
-        return rule_based_convert(text) + "\n\n[Примечание: используется правило-базовое преобразование (ZeroGPU отключён).]"
-
+    # ZeroGPU-only path; if it fails, show an informative message (no rule-based output)
+    try:
+        return _infer_zerogpu(prompt, gen_kwargs)
+    except Exception as e:
+        return f"[Ошибка ZeroGPU: {type(e).__name__}: {e}]"
 
 # ========= UI =========
 with gr.Blocks(title="Pre-reform → Modern Russian (ZeroGPU)") as demo:
     gr.Markdown(
         """
         # Преобразование дореформенной → современной орфографии
-        По умолчанию генерация выполняется на **ZeroGPU** (GPU выделяется на время запроса).  
-        Если ZeroGPU временно недоступен, используется надёжный **правило-базовый** конвертер.
+        Запросы выполняются на **ZeroGPU** (GPU выделяется только на время генерации).
         """
     )
 
@@ -153,18 +148,26 @@ with gr.Blocks(title="Pre-reform → Modern Russian (ZeroGPU)") as demo:
                 repetition_penalty = gr.Slider(1.0, 2.0, value=1.05, step=0.01, label="repetition_penalty")
             btn = gr.Button("Преобразовать", variant="primary")
         with gr.Column():
-            out = gr.Textbox(label="Вывод: современная орфография", lines=12)
+            out = gr.Textbox(label="Вывод: современная орфография", lines=14)
 
     gr.Examples(
         examples=[
+            # Classic prose examples
             ["въ семъ домѣ обитало три семейства, и каждое имѣло свои обыкновенія."],
-            ["Онъ шёлъ по узкой улѣцѣ, разсматривая вывѣски лавокъ и фонари."]
+            ["Онъ шёлъ по узкой улѣцѣ, разсматривая вывѣски лавокъ и фонари."],
+            ["въ мирѣ сёмъ многа есть, чего мудрецу и не снилось."],
+            # Orthography stress tests
+            ["Сей образъ мыслей былъ въ обычаѣ: въслѣдствіе того, что ѣще не наступило прояснѣніе."],
+            ["Именіе его находилось на уѣздной окраинѣ; крестьяне имѣли обыкновеніе собираться къ вечеру."],
+            ["Лѣтописи глаголютъ, яко многа бывало чудесъ на рѣкѣ сей."],
+            ["Оный человѣкъ писалъ послѣднія строки при свѣтѣ фонаря, на улицѣ безлюдной."],
+            ["Въ семъ письмѣ обрѣтёте вы извѣстія, коихъ до нынѣ не имѣли."],
         ],
         inputs=[inp],
     )
 
     btn.click(
-        lambda t,a,b,c,d,e: convert(t, a, b, c, d, e, False),
+        lambda t,a,b,c,d,e: convert(t, a, b, c, d, e),
         inputs=[inp, max_new_tokens, temperature, top_p, top_k, repetition_penalty],
         outputs=[out],
     )