import streamlit as st
from transformers import pipeline
import torch

translator = pipeline("translation",
                      model="facebook/nllb-200-distilled-600M",
                      torch_dtype=torch.bfloat16)

#lang = ["Yoruba","Afrikaans", "English"]

# src = st.selectbox(
#            "Select source language",
#            options=lang
#         )

# tgt = st.selectbox(
#            "Select target language",
#            options=lang
#         )

african_languages = { "English" : "eng_Latn",
"Afrikaans": "afr_Latn",
"Bambara": "bam_Latn",
"Basque": "eus_Latn",
"Ewe": "ewe_Latn",
"Fon": "fon_Latn",
"Fula": "fuv_Latn",
"Hausa": "hau_Latn",
"Igbo": "ibo_Latn",
"Kabyle": "kab_Latn",
"Kalenjin": "kal_Latn",
"Kamba": "kam_Latn",
"Kanjobal": "ktz_Latn",
"Kikuyu": "kik_Latn",
"Kinyarwanda": "kin_Latn",
"Komi": "kom_Cyrl",
"Kongo": "kon_Latn",
"Lingala": "lin_Latn",
"Lozi": "loz_Latn",
"Lubumbashi Swahili": "lui_Latn",
"Luo": "luo_Latn",
"Machame Chaga": "cgc_Latn",
"Madurese": "mad_Latn",
"Maore Comorian": "com_Arab",
"Meru": "mer_Latn",
"Mozambican Portuguese": "pcm_Latn",
"Ngomba": "gnb_Latn",
"Nigerian Fulfude": "fuv_Latn",
"North Levelnde Arabic": "apc_Arab",
"Northern Lwo": "lwo_Latn",
"Northern Ndebele": "lin_Latn",
"Northern Sotho": "nso_Latn",
"Northern Wolof": "wol_Latn",
"Nyanja": "nya_Latn",
"Omotic - Benchnon": "bbg_Latn",
"Oromo": "gaz_Latn",
"Piedmontese": "pms_Latn",
"Rombo": "rmb_Latn",
"Runyankore": "run_Latn",
"Rwa": "rwk_Latn",
"Sena": "seh_Latn",
"Setswana": "tsn_Latn",
"Shona": "sna_Latn",
"Silozi": "loz_Latn",
"Somali": "som_Latn",
"Southern Ndebele": "nbl_Latn",
"Southern Sotho": "sot_Latn",
"Sudanese Colloquial Arabic": "apd_Arab",
"Swazi": "ssw_Latn",
"Swahili": "swh_Latn",
"Tachelhit": "shi_Tfng",
"Taita": "dtp_Latn",
"Tamasheq (Latin script)": "taq_Latn",
"Tigrigna": "tir_Ethi",
"Tswana": "tsn_Latn",
"Tumbuka": "tum_Latn",
"Tutong Murut": "ttm_Latn",
"Twii": "twi_Latn",
"Umbundu": "umb_Latn",
"Venda": "ven_Latn",
"Walloon": "wln_Latn",
"Wolof": "wol_Latn",
"Xitsonga": "tso_Latn",
"Yoruba": "yor_Latn",
"Zaramo": "zmw_Latn",
"Zulu": "zul_Latn",
}

lang = list(african_languages.keys())


st.header('African Languages Translator', divider='rainbow')



#input_text = st.text_input("Enter text to be translated")
#translated_text = translator(input_text, src_lang= lang_dict[src], tgt_lang= lang_dict[tgt])[0]["translation_text"]

translated_text = translator(st.text_input("Enter text to be translated"), src_lang= african_languages[st.selectbox(
            "Select source language",
            options=lang
      )], tgt_lang= african_languages[st.selectbox(
            "Select target language",
            options=lang
       )])[0]["translation_text"]

if st.button('Translate'):
    st.write("Translated Text:", translated_text)


# if input_text and src and tgt and st.button('Translate'):
    
#     st.write("Translated Text:", translated_text)
#     #st.text(translated_text)