import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

# 1. Указываем специализированную чеченскую модель
model_name = "NM-development/nllb-ce-rus-v0"

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

def translate_ru_to_ce(text):
    # Для этой модели НЕ НУЖНО вручную прописывать src_lang и forced_bos_token_id
    inputs = tokenizer(text, return_tensors="pt", padding=True)
    
    # Генерируем перевод напрямую
    outputs = model.generate(**inputs)
    
    # Декодируем результат (outputs[0] берет первую строку из батча)
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

# Интерфейс Gradio
interface = gr.Interface(
    fn=translate_ru_to_ce, 
    inputs=gr.Textbox(label="Введите текст на русском"), 
    outputs=gr.Textbox(label="Перевод на чеченский")
)

interface.launch()