import streamlit as st
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
import torch

# Optionally display the torch version for verification
st.sidebar.write(f'PyTorch Version: {torch.__version__}')

# Load the model and tokenizer for Georgian to English translation
model_name = "Helsinki-NLP/opus-mt-ka-en"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)

# UI Title
st.title("Georgian to English Translator")

# UI for input text
text = st.text_area("Enter Georgian text:", height=150)

if st.button('Translate'):
    if text:
        # Prepare the text for translation
        inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512)
        # Translate the text
        translated = model.generate(**inputs)
        # Decode the translated text
        translation = tokenizer.decode(translated[0], skip_special_tokens=True)
        # Display the translation
        st.write(translation)
    else:
        st.write("Please enter some text to translate.")