Update app.py

app.py

@@ -1,167 +1,208 @@
 import os
-import random
-import uuid
-import json
-import time
-import asyncio
-from threading import Thread
+import re
+import importlib.util
+from pathlib import Path
 
-import gradio as gr
-import spaces
 import torch
-import numpy as np
-from PIL import Image, ImageOps
-# import cv2  # not needed anymore
+import gradio as gr
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
+
+MODEL_ID = "ZennyKenny/oss-20b-prereform-to-modern-ru-merged"
+
+# ----------------- Load SYSTEM_PROMPT from 'text-prompt.py' -----------------
+def _load_system_prompt():
+    prompt_path = Path(__file__).with_name("text-prompt.py")
+    default_prompt = (
+        "Ты компетентный редактор русского языка. "
+        "Преобразуй дореформенную русскую орфографию (до 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)
+    except Exception:
+        return default_prompt
+
+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
+
+# ----------------- 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"
+        f"Текст (дореформ.):\n{text.strip()}\n\n"
+        f"Текст (современная орфография):"
+    )
 
-from transformers import (
-    Qwen2_5_VLForConditionalGeneration,
-    AutoProcessor,
-    TextIteratorStreamer,
-)
-from transformers.image_utils import load_image
+def load_model_cpu():
+    """Force CPU load. Gracefully degrade if loading fails."""
+    global _tokenizer, _model, _streamer, _MODEL_READY, _MODEL_ERROR
+    if _MODEL_READY or _MODEL_ERROR:
+        return
 
-# Optional docling imports (unused now but kept for easy re-enable)
-# from docling_core.types.doc import DoclingDocument, DocTagsDocument
+    if os.getenv("DISABLE_MODEL", "0") == "1":
+        _MODEL_ERROR = "Model disabled via DISABLE_MODEL=1."
+        return
 
-import re
-import ast
-import html
-
-# ---------------------------
-# Constants & device
-# ---------------------------
-MAX_MAX_NEW_TOKENS     = 2048
-DEFAULT_MAX_NEW_TOKENS = 1024
-MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096"))
-
-device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-
-# ---------------------------
-# Load ONLY Typhoon OCR 20B
-# ---------------------------
-MODEL_ID = "scb10x/typhoon-ocr-20b"  # <- 20B model
-processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
-model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
-    MODEL_ID,
-    trust_remote_code=True,
-    torch_dtype=torch.float16
-).to(device).eval()
-
-# ---------------------------
-# (Optional) image helpers
-# ---------------------------
-def add_random_padding(image, min_percent=0.1, max_percent=0.10):
-    image = image.convert("RGB")
-    width, height = image.size
-    pad_w_percent = random.uniform(min_percent, max_percent)
-    pad_h_percent = random.uniform(min_percent, max_percent)
-    pad_w = int(width * pad_w_percent)
-    pad_h = int(height * pad_h_percent)
-    corner_pixel = image.getpixel((0, 0))
-    padded_image = ImageOps.expand(image, border=(pad_w, pad_h, pad_w, pad_h), fill=corner_pixel)
-    return padded_image
-
-def normalize_values(text, target_max=500):
-    def normalize_list(values):
-        max_value = max(values) if values else 1
-        return [round((v / max_value) * target_max) for v in values]
-
-    def process_match(match):
-        num_list = ast.literal_eval(match.group(0))
-        normalized = normalize_list(num_list)
-        return "".join([f"<loc_{num}>" for num in normalized])
-
-    pattern = r"\[([\d\.\s,]+)\]"
-    return re.sub(pattern, process_match, text)
-
-# ---------------------------
-# Image generation only
-# ---------------------------
-@spaces.GPU
-def generate_image(
+    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,
+            trust_remote_code=True,
+            torch_dtype=torch.float32,      # CPU dtype
+            low_cpu_mem_usage=True,
+            device_map=None,                # ensure CPU
+        ).to("cpu")
+        _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,
-    image: Image.Image,
-    max_new_tokens: int = 2048,
-    temperature: float = 0.1,
-    top_p: float = 0.9,
-    top_k: int = 50,
-    repetition_penalty: float = 1.2,
-):
-    """Generate OCR/vision response for a single image with Typhoon OCR 20B."""
-    if image is None:
-        yield "Please upload an image."
-        return
+    max_new_tokens: int,
+    temperature: float,
+    top_p: float,
+    top_k: int,
+    repetition_penalty: float,
+    do_stream: bool
+) -> str:
+    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),
+        top_p=float(top_p),
+        top_k=int(top_k),
+        repetition_penalty=float(repetition_penalty),
+        do_sample=True,
+    )
 
-    images = [image]
+    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)
+        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)
+
+    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
+            )
+        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
+
+# ----------------- UI -----------------
+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 загрузка большой модели может быть недоступна; в таком случае
+        автоматически используется быстрый правило-базовый конвертер (ѣ→е, і→и, ѳ→ф, ѵ→и, удаление конечного ъ).
+        """
+    )
 
-    messages = [
-        {
-            "role": "user",
-            "content": [{"type": "image"} for _ in images] + [
-                {"type": "text", "text": text}
-            ]
-        }
+    with gr.Row():
+        with gr.Column(scale=1):
+            inp = gr.Textbox(
+                label="Ввод: дореформенный текст",
+                placeholder="Например: \"въ мирѣ сёмъ многа есть...\"",
+                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")
+                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(scale=1):
+            out = gr.Textbox(label="Вывод: современная орфография", lines=12)
+
+    examples = [
+        ["въ семъ домѣ обитало три семейства, и каждое имѣло свои обыкновенія."],
+        ["Онъ шёлъ по узкой улѣцѣ, разсматривая вывѣски лавокъ и фонари."],
+        ["въ мирѣ сёмъ многа есть, чего мудрецу и не снилось."]
     ]
-    prompt = processor.apply_chat_template(messages, add_generation_prompt=True)
-    inputs = processor(text=prompt, images=images, return_tensors="pt").to(device)
-
-    streamer = TextIteratorStreamer(processor, skip_prompt=True, skip_special_tokens=True)
-    generation_kwargs = {
-        **inputs,
-        "streamer": streamer,
-        "max_new_tokens": max_new_tokens,
-        "temperature": temperature,
-        "top_p": top_p,
-        "top_k": top_k,
-        "repetition_penalty": repetition_penalty,
-    }
-
-    thread = Thread(target=model.generate, kwargs=generation_kwargs)
-    thread.start()
-
-    buffer = ""
-    for new_text in streamer:
-        buffer += new_text.replace("<|im_end|>", "")
-        yield buffer
-
-# ---------------------------
-# Minimal UI (Image only)
-# ---------------------------
-css = """
-.submit-btn {
-    background-color: #2980b9 !important;
-    color: white !important;
-}
-.submit-btn:hover {
-    background-color: #3498db !important;
-}
-"""
-
-with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
-    gr.Markdown("# **Typhoon OCR 20B**")
+    gr.Examples(examples=examples, inputs=[inp])
 
-    with gr.Row():
-        with gr.Column():
-            image_query = gr.Textbox(label="Query Input", placeholder="e.g., \"OCR the image\" or task instruction…")
-            image_upload = gr.Image(type="pil", label="Image")
-            image_submit = gr.Button("Submit", elem_classes="submit-btn")
-
-            with gr.Accordion("Advanced options", open=False):
-                max_new_tokens = gr.Slider(label="Max new tokens", minimum=1, maximum=MAX_MAX_NEW_TOKENS, step=1, value=DEFAULT_MAX_NEW_TOKENS)
-                temperature = gr.Slider(label="Temperature", minimum=0.1, maximum=4.0, step=0.1, value=0.1)
-                top_p = gr.Slider(label="Top-p (nucleus sampling)", minimum=0.05, maximum=1.0, step=0.05, value=0.9)
-                top_k = gr.Slider(label="Top-k", minimum=1, maximum=1000, step=1, value=50)
-                repetition_penalty = gr.Slider(label="Repetition penalty", minimum=1.0, maximum=2.0, step=0.05, value=1.2)
-
-        # Right column: ONLY output (no model info, no radios)
-        with gr.Column():
-            output = gr.Textbox(label="Output", interactive=False, lines=12, scale=2)
-
-    image_submit.click(
-        fn=generate_image,
-        inputs=[image_query, image_upload, max_new_tokens, temperature, top_p, top_k, repetition_penalty],
-        outputs=output
+    def _on_click(text, a, b, c, d, e, f):
+        return convert(text, a, b, c, d, e, f)
+
+    btn.click(
+        _on_click,
+        inputs=[inp, max_new_tokens, temperature, top_p, top_k, repetition_penalty, do_stream],
+        outputs=[out]
     )
 
 if __name__ == "__main__":
-    demo.queue(max_size=30).launch(share=True, mcp_server=True, ssr_mode=False, show_error=True)
+    os.environ.setdefault("HF_HUB_DISABLE_TELEMETRY", "1")
+    demo.queue().launch()