import streamlit as st
from transformers import MarianMTModel, MarianTokenizer

# Path to your model files
model_path = "lz039/translation-ft-de-bg"
tokenizer = MarianTokenizer.from_pretrained(model_path)
model = MarianMTModel.from_pretrained(model_path)

# Check if loaded successfully
print("Model and tokenizer loaded!")

# Streamlit app
st.title("Translation App")
st.write("Translate text from German to Bulgarian using your custom model!")

# Input text box
src_text = st.text_area("Enter text to translate:", placeholder="Type text in German here...")

# Button to trigger translation
if st.button("Translate"):
    if src_text.strip():  # Ensure there's input text
        try:
            # Tokenize and generate translation
            inputs = tokenizer(src_text, return_tensors="pt", padding=True)
            translated = model.generate(**inputs)
            
            # Decode the translation
            result = tokenizer.decode(translated[0], skip_special_tokens=True)
            
            # Display the translation
            st.success("Translation:")
            st.write(result)
        except Exception as e:
            st.error(f"An error occurred: {e}")
    else:
        st.warning("Please enter some text to translate!")