LingoLink / app.py
shubhendu-ghosh's picture
Update app.py
2cc2cb2 verified
import streamlit as st
import requests
import os
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
from transformers import pipeline
st.markdown('<p style="text-align:center; font-size: 70px; margin-bottom: 50px; margin-top: 0px;"><b>LingoLink</b></p>', unsafe_allow_html=True)
# 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"]
lang_dict = {
"Assamese":"asm_Beng",
"Bengali":"ben_Beng",
"Chinese":"zho_Hans",
"Czech":"ces_Latn",
"Danish":"dan_Latn",
"English":"eng_Latn",
"German":"deu_Latn",
"Greek":"ell_Grek",
"Finnish":"fin_Latn",
"French":"fra_Latn",
"Gujarati":"guj_Gujr",
"Hindi":"hin_Deva",
"Japanese":"jpn_Jpan",
"Kannada":"kan_Knda",
"Korean":"kor_Hang",
"Maithili":"mai_Deva",
"Malayalam":"mal_Mlym",
"Odia":"ory_Orya",
"Portuguese":"por_Latn",
"Russian":"rus_Cyrl",
"Sanskrit":"san_Deva",
"Tamil":"tam_Taml",
"Telugu":"tel_Telu",
"Urdu":"urd_Arab",
}
languages = tuple(lang_dict.keys())
def main():
cols=st.columns(2)
with cols[0]:
src_lang = st.selectbox('From', languages)
with cols[1]:
tgt_lang = st.selectbox('To', languages)
src_lang = lang_dict[src_lang]
tgt_lang = lang_dict[tgt_lang]
text = st.text_area("Enter your text below:")
if st.button("Translate"):
if src_lang and tgt_lang and text:
result = translate(text, src_lang, tgt_lang)
st.text_area("", result)
else:
st.warning("Please provide source language, target language, and text to translate.")
if __name__ == "__main__":
main()