app.py

import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, BitsAndBytesConfig
from peft import PeftModel
import tempfile
import os

# Create a temporary directory for offloading
offload_dir = tempfile.mkdtemp()

# Configure 4-bit quantization for efficiency
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.bfloat16,
    bnb_4bit_quant_type="nf4",
)

# Load base model
base_model_id = "facebook/nllb-200-distilled-1.3B"
tokenizer = AutoTokenizer.from_pretrained(base_model_id)
base_model = AutoModelForSeq2SeqLM.from_pretrained(
    base_model_id,
    quantization_config=bnb_config,
    device_map="auto",
    offload_folder=offload_dir  # Add offload directory
)

# Load LoRA adapter
model = PeftModel.from_pretrained(
    base_model, 
    "Youme21/Manhwa-NLLB-1.3B-EnAr-v2",
    offload_folder=offload_dir  # Add offload directory here too
)
model.eval()

def translate_text(text):
    if not text:
        return ""
    
    inputs = tokenizer(text, return_tensors="pt").to(model.device)
    
    translated_tokens = model.generate(
        **inputs,
        forced_bos_token_id=tokenizer.convert_tokens_to_ids("arb_Arab"),
        max_length=150,
        num_beams=5,
        early_stopping=True
    )
    result = tokenizer.batch_decode(translated_tokens, skip_special_tokens=True)[0]
    return result

description_text = """### 🥋 English to Arabic Manhwa Translator
This model is fine-tuned to handle the specific terminology of **Manhwa, Manga, and Webnovels**.

It understands context like:
*   **Murim:** Qi Deviation, Sect Leader, Young Master.
*   **Systems:** [System], Level Up, Quest.
*   **Honorifics:** Oppa, Hyung, Noona."""

demo = gr.Interface(
    fn=translate_text,
    inputs=gr.Textbox(lines=5, placeholder="Enter English text here... (e.g., 'How dare this Young Master look down on me?')"),
    outputs=gr.Textbox(label="Arabic Translation"),
    title="Manhwa Translator (NLLB-1.3B)",
    description=description_text,
    examples=[
        ["The Young Master is suffering from Qi Deviation!"],
        ["[System] You have leveled up."],
        ["Don't look down on me just because I'm a commoner, you trash!"],
        ["Oppa, please buy me some food."]
    ],
    theme="soft"
)

demo.launch()