import torch
import threading
from transformers import (
    AutoTokenizer,
    AutoModelForSeq2SeqLM,
    TextIteratorStreamer,
)

MODEL_NAME = "facebook/nllb-200-distilled-600M"

print("Loading tokenizer and model...")

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)

device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)
model.eval()

# -------- normal (non-streaming) ----------
def translate(text: str):
    inputs = tokenizer(text, return_tensors="pt").to(device)

    lang_id = tokenizer.convert_tokens_to_ids("arb_Arab")

    with torch.no_grad():
        outputs = model.generate(
            inputs["input_ids"],
            forced_bos_token_id=lang_id,
            max_length=300,
        )

    return tokenizer.decode(outputs[0], skip_special_tokens=True)


# -------- streaming ----------
def stream_translate(text: str):
    inputs = tokenizer(text, return_tensors="pt").to(device)

    lang_id = tokenizer.convert_tokens_to_ids("arb_Arab")

    streamer = TextIteratorStreamer(
        tokenizer,
        skip_special_tokens=True,
        skip_prompt=True,
    )

    generation_kwargs = dict(
        input_ids=inputs["input_ids"],
        forced_bos_token_id=lang_id,
        max_length=300,
        streamer=streamer,
    )

    # Run generation in background thread
    thread = threading.Thread(
        target=model.generate,
        kwargs=generation_kwargs,
    )
    thread.start()

    # Yield tokens as they are generated
    for token in streamer:
        yield token