Update app.py

app.py

@@ -2,83 +2,111 @@ import os
 import re
 from pathlib import Path
 
-import torch
 import gradio as gr
+import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
-from peft import PeftModel  # NEW
+from peft import PeftModel
+import spaces  # ZeroGPU
+
 
-MODEL_ID_BASE = "openai/gpt-oss-20b"  # base model
+# ========= 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"   # where adapter lives in your repo
+ADAPTER_SUBFOLDER = "checkpoint-60"   # LoRA lives here in your repo
 
-# ---- load SYSTEM_PROMPT from text-prompt.py (same as before) ----
+# ZeroGPU toggle (you can also set in Space Secrets):
+USE_ZEROGPU = os.getenv("USE_ZEROGPU", "1") == "1"
+
+# ========= Load external prompt =========
 def _load_system_prompt():
-    prompt_path = Path(__file__).with_name("text-prompt.py")
-    default = ("Ты компетентный редактор русского языка. "
-               "Преобразуй дореформенную русскую орфографию (до 1918 года) "
-               "в современную орфографию. Сохраняй смысл, пунктуацию и регистр. "
-               "Не добавляй комментариев. Верни только преобразованный текст.")
+    path = Path(__file__).with_name("text-prompt.py")
+    default = (
+        "Ты компетентный редактор русского языка. "
+        "Преобразуй дореформенную русскую орфографию (до 1918 года) "
+        "в современную орфографию. Сохраняй смысл, пунктуацию и регистр. "
+        "Не добавляй комментариев. Верни только преобразованный текст."
+    )
     try:
         ns = {}
-        exec(prompt_path.read_text(encoding="utf-8"), ns) if prompt_path.exists() else None
+        if path.exists():
+            exec(path.read_text(encoding="utf-8"), ns)
         return ns.get("SYSTEM_PROMPT", default)
     except Exception:
         return default
 
 SYSTEM_PROMPT = _load_system_prompt()
 
-# ---- simple rule-based fallback (unchanged) ----
-REPLACEMENTS = [("Ѣ","Е"),("ѣ","е"),("І","И"),("і","и"),("Ѳ","Ф"),("ѳ","ф"),("Ѵ","И"),("ѵ","и")]
+def build_prompt(text: str) -> str:
+    return (
+        f"{SYSTEM_PROMPT}\n\n"
+        f"Текст (дореформ.):\n{text.strip()}\n\n"
+        f"Текст (современная орфография):"
+    )
+
+# ========= Rule-based CPU fallback =========
+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) ----
-_tokenizer = None
-_model = None
-_streamer = None
-_MODEL_READY = False
-_MODEL_ERROR = None
 
-def build_prompt(text: str) -> str:
-    return f"{SYSTEM_PROMPT}\n\nТекст (дореформ.):\n{text.strip()}\n\nТекст (современная орфография):"
-
-def load_model_cpu():
-    """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
+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)
+def _infer_zerogpu(prompt: str, gen_kwargs: dict) -> str:
+    # Load tokenizer from your 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.
+    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,
+        trust_remote_code=True,
+        torch_dtype=torch_dtype,
+        device_map="auto",
+    )
+
+    # 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
     try:
-        os.environ["CUDA_VISIBLE_DEVICES"] = ""
-        _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,
-            low_cpu_mem_usage=True,
-            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, max_new_tokens, temperature, top_p, top_k, repetition_penalty, do_stream):
+        model = model.merge_and_unload()
+    except Exception:
+        pass
+
+    # Generate on GPU
+    input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(model.device)
+    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)
+
+    marker = "Текст (современная орфография):"
+    return out.split(marker, 1)[-1].strip() if marker in out else out.strip()
+
+
+# ========= Orchestrator =========
+def convert(text, max_new_tokens, temperature, top_p, top_k, repetition_penalty, do_stream):
+    if not text or not text.strip():
+        return ""
+
     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),
@@ -87,64 +115,57 @@ def convert_with_model(text, max_new_tokens, temperature, top_p, top_k, repetiti
         repetition_penalty=float(repetition_penalty),
         do_sample=True,
     )
-    if do_stream:
-        chunks = []
-        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():
-            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:
+    # Prefer ZeroGPU if enabled; otherwise CPU fallback
+    if USE_ZEROGPU:
         try:
-            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.]"
-    note = "\n\n[Примечание: используется правило-базовое преобразование"
-    if _MODEL_ERROR: note += f" (модель недоступна: {_MODEL_ERROR})"
-    note += ".]"
-    return rule_based_convert(text) + note
-
-# ---- Gradio UI (same structure as before) ----
-with gr.Blocks(title="Pre-reform → Modern Russian (CPU-only)") as demo:
+            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 отключён).]"
+
+
+# ========= UI =========
+with gr.Blocks(title="Pre-reform → Modern Russian (ZeroGPU)") as demo:
     gr.Markdown(
         """
-        # Преобразование дореформенной орфографии → современная (CPU-only)
-        Модель: LoRA-адаптер к `openai/gpt-oss-20b` из `ZennyKenny/oss-20b-prereform-to-modern-ru-merged`.
-        При недоступности модели используется правило-базовый конвертер (ѣ→е, і→и, ѳ→ф, ѵ→и, удаление конечного ъ).
+        # Преобразование дореформенной → современной орфографии
+        По умолчанию генерация выполняется на **ZeroGPU** (GPU выделяется на время запроса).  
+        Если ZeroGPU временно недоступен, используется надёжный **правило-базовый** конвертер.
         """
     )
+
     with gr.Row():
         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")
+            inp = gr.Textbox(
+                label="Ввод: дореформенный текст",
+                placeholder="Например: \"въ мирѣ сёмъ многа есть...\"",
+                lines=10
+            )
+            with gr.Accordion("Параметры генерации", open=False):
+                max_new_tokens = gr.Slider(16, 512, value=192, step=8, label="max_new_tokens")
+                temperature = gr.Slider(0.0, 1.0, value=0.2, step=0.05, label="temperature")
                 top_p = gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="top_p")
                 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():
             out = gr.Textbox(label="Вывод: современная орфография", lines=12)
+
     gr.Examples(
-        examples=[["въ семъ домѣ обитало три семейства, и каждое имѣло свои обыкновенія."]],
+        examples=[
+            ["въ семъ домѣ обитало три семейства, и каждое имѣло свои обыкновенія."],
+            ["Онъ шёлъ по узкой улѣцѣ, разсматривая вывѣски лавокъ и фонари."]
+        ],
         inputs=[inp],
     )
+
     btn.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],
+        lambda t,a,b,c,d,e: convert(t, a, b, c, d, e, False),
+        inputs=[inp, max_new_tokens, temperature, top_p, top_k, repetition_penalty],
         outputs=[out],
     )