Update app.py

app.py

@@ -1,20 +1,37 @@
 import os
 from pathlib import Path
+from typing import Optional, Tuple
 
 import gradio as gr
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import (
+    AutoTokenizer,
+    AutoModelForCausalLM,
+    AutoProcessor,
+    Qwen2_5_VLForConditionalGeneration,
+    pipeline,
+)
 from peft import PeftModel
 import spaces  # ZeroGPU
 
 
 # ========= Config =========
-# 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
+ADAPTER_SUBFOLDER = os.getenv("ADAPTER_SUBFOLDER", "checkpoint-60")
 
-# ========= Load external system prompt =========
+OCR_MODEL_ID = os.getenv("OCR_MODEL_ID", "ChatDOC/OCRFlux-3B")
+
+OCR_MAX_NEW_TOKENS = int(os.getenv("OCR_MAX_NEW_TOKENS", "6000"))
+CONVERT_MAX_NEW_TOKENS = int(os.getenv("CONVERT_MAX_NEW_TOKENS", "6000"))
+
+TEMPERATURE = float(os.getenv("CONVERT_TEMPERATURE", "0.2"))
+TOP_P = float(os.getenv("CONVERT_TOP_P", "0.9"))
+TOP_K = int(os.getenv("CONVERT_TOP_K", "40"))
+REPETITION_PENALTY = float(os.getenv("CONVERT_REP_PENALTY", "1.05"))
+
+
+# ========= Load prompts =========
 def _load_system_prompt():
     path = Path(__file__).with_name("text-prompt.py")
     default = (
@@ -26,31 +43,85 @@ 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
 
 SYSTEM_PROMPT = _load_system_prompt()
 
-def build_prompt(text: str) -> str:
+# OCR prompt in its own file
+def _load_ocr_prompt():
+    path = Path(__file__).with_name("ocr-prompt.py")
+    default = (
+        "Извлеки из изображения весь текст БУКВАЛЬНО и на русском языке. "
+        "Ничего не переводить и не исправлять. "
+        "Сохраняй дореформенную орфографию и специальные символы. "
+        "Верни только чистый текст (plain text)."
+    )
+    try:
+        ns = {}
+        if path.exists():
+            exec(path.read_text(encoding="utf-8"), ns)
+        return ns.get("OCR_PROMPT", default)
+    except Exception:
+        return default
+
+OCR_PROMPT = _load_ocr_prompt()
+
+
+def build_conversion_prompt(pre_reform_text: str) -> str:
     return (
         f"{SYSTEM_PROMPT}\n\n"
-        f"Текст (дореформ.):\n{text.strip()}\n\n"
+        f"Текст (дореформ.):\n{pre_reform_text.strip()}\n\n"
         f"Текст (современная орфография):"
     )
 
-# ========= ZeroGPU inference =========
-@spaces.GPU(duration=180)  # GPU is leased only while this function runs
-def _infer_zerogpu(prompt: str, gen_kwargs: dict) -> str:
-    # Tokenizer from adapter repo (it contains tokenizer files)
-    tokenizer = AutoTokenizer.from_pretrained(ADAPTER_REPO, use_fast=True, trust_remote_code=True)
 
-    # Ensure pad token exists; if not, align it with EOS (common for GPT-like)
+# ========= ZeroGPU: OCR step =========
+@spaces.GPU(duration=300)  # 5 minutes
+def _ocr_image_to_text(image) -> str:
+    processor = AutoProcessor.from_pretrained(OCR_MODEL_ID, trust_remote_code=True)
+
+    torch_dtype = torch.bfloat16 if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else torch.float16
+    ocr_model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+        OCR_MODEL_ID,
+        trust_remote_code=True,
+        torch_dtype=torch_dtype,
+        device_map="auto",
+    )
+
+    ocr_pipe = pipeline(
+        task="image-text-to-text",
+        model=ocr_model,
+        processor=processor,
+    )
+
+    out = ocr_pipe(
+        image,
+        prompt=OCR_PROMPT,
+        max_new_tokens=OCR_MAX_NEW_TOKENS,
+        temperature=0.0,
+        do_sample=False,
+    )
+
+    if isinstance(out, list) and len(out) > 0:
+        text = out[0].get("generated_text", "") or out[0].get("text", "")
+    elif isinstance(out, str):
+        text = out
+    else:
+        text = ""
+
+    return (text or "").strip()
+
+
+# ========= ZeroGPU: Conversion step =========
+@spaces.GPU(duration=300)  # 5 minutes
+def _convert_text_zerogpu(pre_reform_text: str) -> str:
+    tokenizer = AutoTokenizer.from_pretrained(ADAPTER_REPO, use_fast=True, trust_remote_code=True)
     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,
@@ -59,43 +130,42 @@ def _infer_zerogpu(prompt: str, gen_kwargs: dict) -> str:
         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
     try:
         model = model.merge_and_unload()
     except Exception:
         pass
 
-    # 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 -----
+    prompt = build_conversion_prompt(pre_reform_text)
     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)
+    gen_kwargs = dict(
+        max_new_tokens=CONVERT_MAX_NEW_TOKENS,
+        temperature=TEMPERATURE,
+        top_p=TOP_P,
+        top_k=TOP_K,
+        repetition_penalty=REPETITION_PENALTY,
+        do_sample=True,
+        use_cache=True,
+    )
 
-    # ----- Generate -----
     with torch.no_grad():
         out_ids = model.generate(
             input_ids=input_ids,
-            attention_mask=attention_mask,  # Key fix for pad==eos
+            attention_mask=attention_mask,
             **gen_kwargs,
         )
 
-    # 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 = "Текст (современная орфография):"
@@ -103,73 +173,56 @@ def _infer_zerogpu(prompt: str, gen_kwargs: dict) -> str:
 
     return out
 
+
 # ========= Orchestrator =========
-def convert(text, max_new_tokens, temperature, top_p, top_k, repetition_penalty):
-    if not text or not text.strip():
-        return ""
+def process(image, manual_text):
+    pre_reform_from_ocr = ""
+    if image is not None:
+        pre_reform_from_ocr = _ocr_image_to_text(image)
 
-    prompt = build_prompt(text)
-    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,
-    )
+    combined = ""
+    if manual_text and manual_text.strip():
+        combined = manual_text.strip()
+    if pre_reform_from_ocr:
+        combined = (combined + "\n\n" + pre_reform_from_ocr).strip() if combined else pre_reform_from_ocr
+
+    if not combined:
+        return "", ""
+
+    modern_text = _convert_text_zerogpu(combined)
+    return modern_text, pre_reform_from_ocr
 
-    # 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:
+with gr.Blocks(title="Pre-reform → Modern Russian (OCR + ZeroGPU)") as demo:
     gr.Markdown(
         """
-        # Преобразование дореформенной → современной орфографии
-        Запросы выполняются на **ZeroGPU** (GPU выделяется только на время генерации).
+        # Преобразование дореформенной → современной орфографии (с OCR)
+        1) Загрузите изображение с дореформенным текстом (фотография/скан), **или** вставьте текст вручную.  
+        2) Модель **OCRFlux-3B** извлечёт текст, затем **OSS-20B + LoRA** преобразует его в современную орфографию.  
+        **Параметры генерации скрыты и настроены для длинных документов (≈ 6 000 токенов).**
         """
     )
 
     with gr.Row():
         with gr.Column():
-            inp = gr.Textbox(
-                label="Ввод: дореформенный текст",
+            img = gr.Image(label="Изображение с дореформенным текстом", type="pil")
+            manual = gr.Textbox(
+                label="(Необязательно) Вставьте дореформенный текст вручную",
+                lines=10,
                 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")
-            btn = gr.Button("Преобразовать", variant="primary")
+            btn = gr.Button("Распознать и преобразовать", variant="primary")
         with gr.Column():
-            out = gr.Textbox(label="Вывод: современная орфография", lines=14)
-
-    gr.Examples(
-        examples=[
-            # Classic prose examples
-            ["въ семъ домѣ обитало три семейства, и каждое имѣло свои обыкновенія."],
-            ["Онъ шёлъ по узкой улѣцѣ, разсматривая вывѣски лавокъ и фонари."],
-            ["въ мирѣ сёмъ многа есть, чего мудрецу и не снилось."],
-            # Orthography stress tests
-            ["Сей образъ мыслей былъ въ обычаѣ: въслѣдствіе того, что ѣще не наступило прояснѣніе."],
-            ["Именіе его находилось на уѣздной окраинѣ; крестьяне имѣли обыкновеніе собираться къ вечеру."],
-            ["Лѣтописи глаголютъ, яко многа бывало чудесъ на рѣкѣ сей."],
-            ["Оный человѣкъ писалъ послѣднія строки при свѣтѣ фонаря, на улицѣ безлюдной."],
-            ["Въ семъ письмѣ обрѣтёте вы извѣстія, коихъ до нынѣ не имѣли."],
-        ],
-        inputs=[inp],
-    )
+            out_modern = gr.Textbox(label="Современная орфография (результат)", lines=18)
+            with gr.Accordion("Промежуточный текст из OCR (для проверки)", open=False):
+                out_ocr = gr.Textbox(label="Текст из OCRFlux-3B", lines=12)
 
     btn.click(
-        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],
+        fn=process,
+        inputs=[img, manual],
+        outputs=[out_modern, out_ocr],
+        api_name="process",
     )
 
 if __name__ == "__main__":