import streamlit as st
import ctranslate2
from transformers import MarianTokenizer

# Load tokenizer and model
tokenizer = MarianTokenizer.from_pretrained("Helsinki-NLP/opus-mt-tr-ar")  # Use the original tokenizer path
translator = ctranslate2.Translator("ctranslate")

st.title("Helsinki-NLP/opus-mt-finetune-tr-ota")

input_text = st.text_area("Enter text to translate:")

if st.button("Translate"):
    # Tokenize input
    tokens = tokenizer.convert_ids_to_tokens(tokenizer.encode(input_text, return_tensors="pt")[0])
    # Translate
    results = translator.translate_batch([tokens], beam_size=5)
    # Detokenize output
    output_tokens = results[0].hypotheses[0]
    translation = tokenizer.decode(tokenizer.convert_tokens_to_ids(output_tokens), skip_special_tokens=True)
    st.success(translation)