import gradio as gr
from transformers import AutoModelForSeq2SeqLM
from transformers import AlbertTokenizer


tokenizer = AlbertTokenizer.from_pretrained(
    "ai4bharat/MultiIndicWikiBioSS", do_lower_case=False, use_fast=False, keep_accents=True)
qgmodel = AutoModelForSeq2SeqLM.from_pretrained(
    "ai4bharat/MultiIndicQuestionGenerationSS").eval()
hgmodel = AutoModelForSeq2SeqLM.from_pretrained(
    "ai4bharat/MultiIndicHeadlineGenerationSS").eval()
ssmodel = AutoModelForSeq2SeqLM.from_pretrained(
    "ai4bharat/MultiIndicSentenceSummarizationSS").eval()
ppmodel = AutoModelForSeq2SeqLM.from_pretrained(
    "ai4bharat/MultiIndicParaphraseGenerationSS").eval()
wbmodel = AutoModelForSeq2SeqLM.from_pretrained(
    "ai4bharat/MultiIndicWikiBioSS").eval()

bos_id = tokenizer._convert_token_to_id_with_added_voc("<s>")
eos_id = tokenizer._convert_token_to_id_with_added_voc("</s>")
pad_id = tokenizer._convert_token_to_id_with_added_voc("<pad>")

INDIC = {"Assamese": "as", "Bengali": "bn", "Gujarati": "gu", "Hindi": "hi", "Kannada": "kn",
         "Malayalam": "ml", "Marathi": "mr", "Odia": "or", "Punjabi": "pa", "Tamil": "ta", "Telugu": "te"}


def generate(input, task, lang):
    lang = INDIC[lang]
    if task == "IndicWikiBio":
        model = wbmodel
    elif task == "IndicHeadlineGeneration":
        model = hgmodel
    elif task == "IndicParaphrasing":
        model = ppmodel
    elif task == "IndicSentenceSummarization":
        model = ssmodel
    elif task == "IndicQuestionGeneration":
        model = qgmodel

    inp = tokenizer(input.strip() + " </s> <2" + lang + ">",
                    add_special_tokens=False, return_tensors="pt", padding=True).input_ids
    model_output = model.generate(inp, use_cache=True, num_beams=1, max_length=100, min_length=1, early_stopping=True, pad_token_id=pad_id,
                                  bos_token_id=bos_id, eos_token_id=eos_id, decoder_start_token_id=tokenizer._convert_token_to_id_with_added_voc("<2"+lang+">"))
    decoded_output = tokenizer.decode(
        model_output[0], skip_special_tokens=True, clean_up_tokenization_spaces=False)

    return decoded_output


tasks = ["IndicWikiBio", "IndicHeadlineGeneration", "IndicParaphrasing",
         "IndicSentenceSummarization", "IndicQuestionGeneration"]
languages = list(INDIC.keys())

language_drop_down = gr.inputs.Dropdown(
    languages, type="value", default="Hindi", label="Select Target Language")
task_drop_down = gr.inputs.Dropdown(
    tasks, type="value", default="IndicSentenceSummarization", label="Select Task")
text = gr.inputs.Textbox(lines=5, placeholder="Enter Indic Text here...",
                         default="", label="Enter Text in corresponding Indic Language")
text_ouptut = gr.outputs.Textbox(
    type="text", label="View Generated Text in the corresponding Indic Language")

supported_lang = ', '.join(languages)

examples = [
    [u"<TAG> name </TAG> राम नरेश पांडेय <TAG> office </TAG> विधायक - 205 - कुशीनगर विधान सभा निर्वाचन क्षेत्र , उत्तर प्रदेश <TAG> term </TAG> 1967 से 1968 <TAG> nationality </TAG> भारतीय", "IndicWikiBio", "Hindi"],
    [u"वैश्विक व्यापार युद्ध की शिकार हुई तुर्की की मुद्रा लीरा के डूबने से अमेरिकी डॉलर के मुकाबले रुपया अब तक के न्यूनतम स्तर पर पहुंच गया। रुपये में रिकॉर्ड गिरावट से सोने की चमक में निखार नहीं आ सकी। वैश्विक बाजार में सोना करीब आठ महीने के निचले स्तर पर पहुंच गया तो घरेलू बाजार में यह करीब नौ महीने के निचले स्तर पर चला गया। वैश्विक मंदी की आशंका से वैश्विक बाजार में चांदी करीब ढाई साल और घरेलू बाजार में तकरीबन नौ महीने के निचले स्तर पर पहुंच गई। तुर्की की आर्थिक चिंता के कारण अमेरिकी डॉलर के मुकाबले रुपया कारोबार के दौरान 70.80 के स्तर तक गिर गया। यह इसका ऐतिहासिक रिकॉर्ड निम्न स्तर है। कमजोर रुपये से सोने की चमक बढऩे की उम्मीद की जा रही थी लेकिन वैश्विक बाजार में सोने की कीमत गिरकर 1,193.50 डॉलर प्रति औंस पहुंचने के कारण घरेलू बाजार में भी सोने की चमक फीकी पड़ गई। घरेलू बाजार में सोना गिरकर 29,655 रुपये प्रति 10 ग्राम पहुंच गया। घरेलू वायदा बाजार यानी एमसीएक्स पर सोना 29,700 के आस-पास कारोबार कर रहा है। देश में इस साल सोने की मांग में लगातार गिरावट देखने को मिल रही थी। अप्रैल-जून तिमाही में सोने का आयात 25 फीसदी से भी कम हुआ है। चालू महीने में सोने की मांग बढऩे की उम्मीद जगी थी लेकिन यह उम्मीद टूट सकती है क्योंकि दुनिया के सबसे बड़े गोल्ड फंड एसपीडीआर गोल्ड की होल्डिंग अप्रैल के बाद 10 फीसदी गिर चुकी है। इस समय यह पिछले ढाई साल के निचले स्तर पर है। इस साल वैश्विक बाजार में सोना करीब 8.5 फीसदी और घरेलू बाजार में 1.5 फीसदी टूट चुका है। सराफा मामलों के जानकार अनिल अग्रवाल कहते हैं कि वैश्विक हालात ऐसे हैं कि इस समय निवेशक डॉलर में पैसा लगा रहे हैं। इस कारण दूसरी मुद्रा और जिंस दबाव में हैं। हालांकि हालात यही रहे तो सोने में तेज सुधार भी देखने को मिलेगा। वैश्विक मंदी की बढ़ती आशंका का सबसे ज्यादा असर चांदी पर पड़ रहा है। वैश्विक बाजार में चांदी के दाम ढाई साल के निचले स्तर पर पहुंच चुके हैं। वैश्विक बाजार में चांदी की कीमत 15 डॉलर प्रति औंस के करीब चल रही है। इसके पहले अप्रैल 2016 में चांदी इस स्तर पर थी। वैश्विक बाजार में चांदी के दाम दो महीने पहले 18.13 डॉलर प्रति औंस पर चल रहे थे। चांदी कारोबारी राहुल मेहता कहते हैं कि सोना और मूल धातु में कमजोरी से चांदी पर दोहरा दबाव पड़ रहा है। वैश्विक बाजार का व्यापार युद्ध अब मुद्रा युद्ध में बदल गया है। वैश्विक अर्थव्यवस्था एक बार फिर मंदी की गिरफ्त में आ सकती है जिसके कारण औद्योगिक विकास भी प्रभावित होगा। यही वजह है कि चांदी की कीमतें लगातार लुढक़ रही हैं क्योंकि मांग में कमी आने की आशंका बढ़ती जा रही है। फिलहाल घरेलू बाजार में चांदी 37,825 रुपये प्रति किलोग्राम पर बिक रही है। तुर्की के आर्थिक संकट से एक बार फिर वैश्विक मंदी का डर है जिसका असर दुनियाभर के बाजारों पर देखा जा सकता है। इसने विश्व स्तर पर निवेशकों के रुख को प्रभावित किया है और वे डॉलर को एक सुरक्षित निवेश के तौर पर देख रहे हैं। आनंद राठी शेयर्स ऐंड स्टाक ब्रोकर्स में शोध विश्लेषक आर मारू ने कहा कि आयातकों की अधिक मांग से रुपये की विनिमय दर में गिरावट आई। उन्होंने कहा, तुर्की संकट को लेकर अनिश्चितता तथा डॉलर सूचकांक में तेजी को देखते हुए आयातक आक्रमक तरीके से डॉलर की लिवाली कर रहे हैं। दूसरी तरफ आरबीआई की तरफ से आक्रमक हस्तक्षेप न होने से भी रुपया नीचे आया। सरकार ने अमेरिकी डॉलर के मुकाबले रुपये के अब तक के न्यूनतम स्तर पर पहुंचने के लिए बाह्य कारकों को जिम्मेदार ठहराते हुए कहा कि इसमें चिंता की कोई बात नहीं है।", "IndicHeadlineGeneration", "Hindi"],
    [u"दिल्ली यूनिवर्सिटी देश की प्रसिद्ध यूनिवर्सिटी में से एक है.",
        "IndicParaphrasing", "Hindi"],
    [u"जम्मू एवं कश्मीर के अनंतनाग जिले में शनिवार को सुरक्षाबलों के साथ मुठभेड़ में दो आतंकवादियों को मार गिराया गया।",
        "IndicSentenceSummarization", "Hindi"],
    [u"7 फरवरी, 2016 [SEP] खेल 7 फरवरी, 2016 को कैलिफोर्निया के सांता क्लारा में सैन फ्रांसिस्को खाड़ी क्षेत्र में लेवी स्टेडियम में खेला गया था।",
        "IndicQuestionGeneration", "Hindi"]
]

iface = gr.Interface(fn=generate, inputs=[text, task_drop_down, language_drop_down], outputs=text_ouptut, title='IndicNLG System',
                     description='Currently the model supports ' + supported_lang, article='More information can be found [here](https://ai4bharat.org/language-generation)', examples=examples)
iface.launch(enable_queue=True)

# with gr.blocks.Blocks() as block:
#     input = gr.Textbox(label="Input")
#     task = gr.Dropdown(["IndicWikiBio", "IndicHeadlineGeneration", "IndicParaphrasing",
#                        "IndicSentenceSummarization", "IndicQuestionGeneration"], label="Task")
#     lang = gr.Dropdown(["as", "bn", "gu", "hi", "kn", "ml",
#                        "mr", "or", "pa", "ta", "te"], label="Language")
#     generate = gr.Button("Generate")
#     output = gr.Textbox()
#     instructions = gr.HTML("<h1>How to use:</h1><br>\
#     1. This space supports 5 tasks and 11 Indic languages.<br>\
#     2. First select the task from the dropdown box and it will show you an example of Input for Hindi. This default example display will be automated for each language in the future. Choose your language, give your input and then press the generate button. Note the formats for IndicWikiBio and Question generation when testing your own inputs. Also note that if you choose another task then the input will be replaced with the default example for that task.<br>\
#     3. The tasks are:<br>\
#     3.1 IndicWikiBio where the input is a Wikipedia table and the output will be a one sentence biograpy. You should pass the input in the following format: &lt;TAG&gt; key1 &lt;/TAG&gt; value1 &lt;TAG&gt; key2 &lt;/TAG&gt; value2.<br>\
#     3.2 IndicHeadlineGeneration where the input is a document or paragraph the output will be a short title. Copy a paragraph from your favorite news site and get a headline. Dont paste extemely long paragraphs. You have been warned.<br>\
#     3.3 IndicParaphrasing where the input is a sentence and the output is its paraphrase.<br>\
#     3.4 IndicSentenceSummarization where the input is a long sentence and the output is a compact version of that sentence.<br>\
#     3.5 IndicQuestionGeneration where the input is an answer and context and the output is the question that should be asked to get the answer. You should pass the input in the following format: ANSWER [SEP] CONTEXT.\
#     ")
#     task.change(set_example, inputs=[task, lang], outputs=[input, lang])
#     generate.click(greet, inputs=[input, task, lang], outputs=output)
#     block.launch()