import streamlit as st
import requests
import os
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
from transformers import pipeline

st.title("Translation App")

# Load the model and tokenizer
tokenizer = AutoTokenizer.from_pretrained("facebook/nllb-200-distilled-600M")
model = AutoModelForSeq2SeqLM.from_pretrained("facebook/nllb-200-distilled-600M")

def translate(text, src_lang, tgt_lang):
    translator = pipeline(
        "translation",
        model=model,
        tokenizer=tokenizer,
        src_lang=src_lang,
        tgt_lang=tgt_lang,
    )
    output = translator(text, max_length=400)
    return output[0]["translation_text"]

def main():
    src_lang = st.text_input("Enter source language code (e.g., en):")
    tgt_lang = st.text_input("Enter target language code (e.g., fr):")
    text = st.text_area("Enter text to translate:")
    
    if st.button("Translate"):
        if src_lang and tgt_lang and text:
            result = translate(text, src_lang, tgt_lang)
            st.write("Translated Text:", result)
        else:
            st.warning("Please provide source language, target language, and text to translate.")

if __name__ == "__main__":
    main()