import os
import asyncio
import torch
import streamlit as st
from transformers import MBartForConditionalGeneration, MBart50TokenizerFast
from transformers import MBart50Tokenizer


os.environ["STREAMLIT_WATCH_FILE_SYSTEM"] = "false"

try:
    asyncio.get_running_loop()
except RuntimeError:
    asyncio.set_event_loop(asyncio.new_event_loop())

# Load model and tokenizer
model_path = "app90/ChhattishgarhiAI_Model"
model = MBartForConditionalGeneration.from_pretrained(model_path)
tokenizer = MBart50Tokenizer.from_pretrained(model_path, src_lang="hi_IN", tgt_lang="hne_IN")
tokenizer.save_pretrained("CG_AI")

# Translate Hindi → Chhattisgarhi
def translate_hindi_to_chhattisgarhi(text):
    sentences = text.split("।")
    translated_sentences = []

    for sentence in sentences:
        sentence = sentence.strip()
        if sentence:
            inputs = tokenizer(sentence, return_tensors="pt", truncation=True, padding="longest", max_length=256)
            with torch.no_grad():
                translated_ids = model.generate(**inputs, max_length=256, num_beams=5, early_stopping=True)
            translated_text = tokenizer.decode(translated_ids[0], skip_special_tokens=True)
            translated_sentences.append(translated_text)

    return " । ".join(translated_sentences)

# Streamlit UI
st.title("Hindi to Chhattisgarhi Translator 🗣️")
st.write("Enter a Hindi sentence and get its translation in Chhattisgarhi.")

user_input = st.text_area("Enter text:")

if st.button("Translate"):
    if user_input.strip():
        chhattisgarhi_text = translate_hindi_to_chhattisgarhi(user_input)
        st.success(f"**Chhattisgarhi Translation**:\n{chhattisgarhi_text}")
    else:
        st.warning("⚠ Please enter some text before translating.")