import streamlit as st
from transformers import MarianMTModel, AutoTokenizer,pipeline
import torch

st.title("Translate Fine tunning model")

# Load model and tokenizer
model_dir = "translate_model_finetune"
tokenizer = AutoTokenizer.from_pretrained(model_dir)
model = MarianMTModel.from_pretrained(model_dir)


def translate(text, src_lang="en", tgt_lang="hi"):
    inputs = tokenizer(text, return_tensors="pt")
    
    translated_tokens = model.generate(**inputs, max_length=50)
    return tokenizer.decode(translated_tokens[0], skip_special_tokens=True)




inputs_text=st.text_input("Please enter the text",value="good boy")

if st.button("submit"):
    output= translate(inputs_text)
    
    st.write(output)