Upload app.py

app.py

@@ -0,0 +1,237 @@
+import os
+import re
+import time
+
+import spaces
+import torch
+import gradio as gr
+import pysubs2
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+# ── تنظیمات ──────────────────────────────────────────────
+MODEL_ID = "google/translategemma-4b-it"
+SRC_LANG = "en"
+TGT_LANG = "fa"
+BATCH_SIZE = 12
+MAX_NEW_TOKENS = 1536
+
+HF_TOKEN = os.environ.get("HF_TOKEN")
+
+# ── بارگذاری مدل (سطح global، طبق الزام ZeroGPU) ────────
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID,
+    torch_dtype=torch.bfloat16,
+    token=HF_TOKEN,
+)
+model.to("cuda")
+model.eval()
+
+
+# ── ساخت پرامپت و پارس خروجی ─────────────────────────────
+def build_batch(subs, start, end):
+    chunks = []
+    for n, idx in enumerate(range(start, end), 1):
+        chunks.append(f"SEGMENT {n}:\n{subs[idx].text.strip()}")
+    return "\n\n".join(chunks)
+
+
+def build_prompt(batch_text, count):
+    return f"""<start_of_turn>user
+شما یک مترجم حرفه‌ای زیرنویس انگلیسی به فارسی هستید.
+
+متن زیر شامل دقیقاً {count} بخش (SEGMENT) پیاپی از یک گفتگوی طبیعی است. هر بخش را با توجه به بافت کل گفتگو، به فارسی محاوره‌ای، خودمونی و روان ترجمه کن — مثل زبان صحبت روزمره، نه فارسی رسمی و کتابی.
+
+قوانین مهم:
+- هر SEGMENT را جداگانه و دقیقاً به همان ترتیب ترجمه کن.
+- هرگز دو یا چند SEGMENT را با هم ادغام نکن و هیچ SEGMENT را نصف نکن.
+- برچسب‌های SEGMENT فقط برای هم‌ترازی هستند؛ آن‌ها را ترجمه نکن.
+- نام‌های خاص، اصطلاحات فنی و کلمات انگلیسی رایج (مثل اسم نرم‌افزارها) را به همون شکل انگلیسی نگه دار.
+- حتی اگر یک SEGMENT وسط جمله (بدون نقطه‌ی پایان) قطع شده باشد، فقط همان کلمات را ترجمه کن؛ هرگز ادامه‌ی جمله یا محتوای بعدی را حدس نزن یا اضافه نکن.
+- فقط خروجی خواسته‌شده را بده، بدون توضیح اضافه.
+
+خروجی را دقیقاً به این شکل بده:
+SEGMENT 1: <ترجمه فارسی>
+SEGMENT 2: <ترجمه فارسی>
+...
+SEGMENT {count}: <ترجمه فارسی>
+
+گفتگو:
+{batch_text}<end_of_turn>
+<start_of_turn>model
+"""
+
+
+def parse_output(text, count):
+    text = (text or "").strip()
+    text = re.sub(r"```(?:text|plaintext|srt)?", "", text, flags=re.I).replace("```", "").strip()
+    text = re.sub(
+        r"(?im)^\s*\[?\s*SEGMENT\s+(\d+)\s*\]?\s*[:：-]\s*",
+        lambda m: f"\n@@SEGMENT_{m.group(1)}@@\n",
+        text
+    )
+    parts = re.split(r"\n@@SEGMENT_(\d+)@@\n", "\n" + text + "\n")
+    found = {}
+    for i in range(1, len(parts), 2):
+        try:
+            n = int(parts[i])
+            value = parts[i + 1].strip()
+            if value:
+                found[n] = value
+        except Exception:
+            pass
+
+    if count == 1 and not found:
+        return [text]
+
+    if len(found) != count or any(i not in found for i in range(1, count + 1)):
+        raise ValueError(f"Expected {count} segments, got {len(found)}")
+    return [found[i] for i in range(1, count + 1)]
+
+
+# ── تابع GPU (هر batch یک فراخوانی کوتاه است) ────────────
+def compute_max_new_tokens(batch_text):
+    # سقف پویا: حدود ۸ توکن به‌ازای هر کلمه‌ی ورودی (کافی برای ترجمه + یک حاشیه‌ی امن)
+    # این کار مانع می‌شود که مدل برای یک تکه‌ی ناقص (بدون نقطه‌ی پایان)
+    # به‌جای ترجمه‌ی همان تکه، ادامه‌ی کل جمله/متن را حدس بزند و تولید کند.
+    word_count = max(1, len(batch_text.split()))
+    return max(48, min(MAX_NEW_TOKENS, word_count * 8))
+
+
+@spaces.GPU(duration=60)
+def translate_batch(batch_text, count):
+    prompt = build_prompt(batch_text, count)
+    inputs = tokenizer(prompt, return_tensors="pt")
+    inputs = {k: v.to(model.device) for k, v in inputs.items()}
+
+    dynamic_max_tokens = compute_max_new_tokens(batch_text)
+
+    with torch.inference_mode():
+        output = model.generate(
+            **inputs,
+            max_new_tokens=dynamic_max_tokens,
+            do_sample=False,
+            pad_token_id=tokenizer.eos_token_id,
+        )
+
+    generated = output[0][inputs["input_ids"].shape[-1]:]
+    answer = tokenizer.decode(generated, skip_special_tokens=True)
+    return parse_output(answer, count)
+
+
+def translate_batch_retry(subs, start, end):
+    batch_text = build_batch(subs, start, end)
+    count = end - start
+    try:
+        return translate_batch(batch_text, count)
+    except Exception as e:
+        print(f"⚠️ Batch alignment failed: {e}")
+        print("🔁 Retrying segment-by-segment...")
+        result = []
+        for idx in range(start, end):
+            try:
+                single = build_batch(subs, idx, idx + 1)
+                result.extend(translate_batch(single, 1))
+            except Exception as e2:
+                print(f"❌ Segment {idx + 1}: {e2}")
+                result.append(subs[idx].text.strip())
+        return result
+
+
+def fix_bidi(text):
+    """دور کلمات/عبارات لاتین و عددی، کاراکتر Unicode Isolate می‌ذاره
+    تا در متن راست‌چین، جهت‌شون به‌هم نریزه."""
+    LRI = "\u2066"  # Left-to-Right Isolate
+    PDI = "\u2069"  # Pop Directional Isolate
+    pattern = re.compile(r'[A-Za-z0-9][A-Za-z0-9 _\-\.]*[A-Za-z0-9]|[A-Za-z0-9]')
+    return pattern.sub(lambda m: f"{LRI}{m.group(0)}{PDI}", text)
+
+
+def translate_srt_file(input_path, progress=gr.Progress()):
+    subs = pysubs2.load(input_path, encoding="utf-8")
+    total = len(subs)
+    translations = [""] * total
+    started = time.time()
+
+    for start in range(0, total, BATCH_SIZE):
+        end = min(total, start + BATCH_SIZE)
+        translations[start:end] = translate_batch_retry(subs, start, end)
+
+        elapsed = time.time() - started
+        print(f"✅ {end}/{total} | {end/elapsed:.2f} segments/sec", flush=True)
+        if progress:
+            progress(end / total, desc=f"ترجمه {end}/{total}")
+
+    output = pysubs2.SSAFile()
+    preview = []
+
+    for i, sub in enumerate(subs):
+        original = sub.text.strip()
+        translated = fix_bidi(translations[i].strip()) if original else ""
+
+        output.append(pysubs2.SSAEvent(
+            start=sub.start,
+            end=sub.end,
+            text=translated
+        ))
+
+        preview.append({
+            "number": i + 1,
+            "english": original,
+            "persian": translated
+        })
+
+    base = os.path.splitext(os.path.basename(input_path))[0]
+    output_path = f"/tmp/{base}_FA.srt"
+    output.save(output_path, encoding="utf-8")
+
+    print(f"✅ Finished in {(time.time()-started)/60:.1f} minutes")
+    return output_path, preview
+
+
+def ui_translate(file, progress=gr.Progress()):
+    if file is None:
+        raise gr.Error("لطفاً یک فایل SRT انتخاب کنید.")
+
+    output_path, data = translate_srt_file(file, progress)
+
+    blocks = []
+    for x in data[:50]:
+        blocks.append(
+            f"### #{x['number']}\n\n"
+            f"🇬🇧 **English**\n\n{x['english']}\n\n"
+            f"🇮🇷 **فارسی**\n\n{x['persian']}\n\n---"
+        )
+    preview = "\n\n".join(blocks)
+    if len(data) > 50:
+        preview += f"\n\n> نمایش ۵۰ مورد اول از {len(data)} subtitle."
+    return output_path, preview
+
+
+with gr.Blocks(title="TranslateGemma 4B SRT Translator") as demo:
+    gr.Markdown("""
+# 🎬 English → Persian SRT Translator
+### Powered by TranslateGemma 4B
+
+مدل timestamp و شماره SRT را نمی‌بیند؛ فقط متن تمیز و segmentهای context-aware را ترجمه می‌کند.
+""")
+
+    with gr.Row():
+        input_file = gr.File(
+            label="📂 English SRT",
+            file_types=[".srt"],
+            type="filepath"
+        )
+        output_file = gr.File(label="📥 Persian SRT")
+
+    translate_button = gr.Button("🚀 شروع ترجمه", variant="primary", size="lg")
+    gr.Markdown("## 👁️ پیش‌نمایش")
+    preview = gr.Markdown()
+
+    translate_button.click(
+        fn=ui_translate,
+        inputs=input_file,
+        outputs=[output_file, preview]
+    )
+
+demo.queue(max_size=5).launch()