import streamlit as st
from transformers import MarianMTModel, MarianTokenizer
import torch

# Load pre-trained models and tokenizers for English to French and English to Hindi
model_name_fr = 'Helsinki-NLP/opus-mt-en-fr'
model_name_hi = 'Helsinki-NLP/opus-mt-en-hi'

tokenizer_fr = MarianTokenizer.from_pretrained(model_name_fr)
model_fr = MarianMTModel.from_pretrained(model_name_fr)

tokenizer_hi = MarianTokenizer.from_pretrained(model_name_hi)
model_hi = MarianMTModel.from_pretrained(model_name_hi)

# Determine the device to use (GPU if available, else CPU)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model_fr.to(device)
model_hi.to(device)

# Function to translate a word
def translate_word(word, tokenizer, model, device):
    inputs = tokenizer.encode(word, return_tensors="pt").to(device)
    outputs = model.generate(inputs)
    translation = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return translation

# Streamlit UI
st.title("10-Letter Word Translator")
st.write("Translate 10-letter English words to French or Hindi")

word = st.text_input("Enter a 10-letter English word")

if word:
    if len(word) != 10:
        st.write("Please enter a 10-letter word.")
    else:
        language = st.selectbox("Select target language", ("French", "Hindi"))
        
        if language == "French":
            translation = translate_word(word, tokenizer_fr, model_fr, device)
        else:
            translation = translate_word(word, tokenizer_hi, model_hi, device)
        
        st.write(f"Translation in {language}: {translation}")