Spaces:
Paused
Paused
| import gradio as gr | |
| import random | |
| import torch | |
| from transformers import MT5Tokenizer, MT5ForConditionalGeneration | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| tokenizer = MT5Tokenizer.from_pretrained("potsawee/mt5-english-thai-large-translation") | |
| translator = MT5ForConditionalGeneration.from_pretrained("potsawee/mt5-english-thai-large-translation") | |
| summarizer = MT5ForConditionalGeneration.from_pretrained("potsawee/mt5-english-thai-large-summarization") | |
| translator.eval() | |
| summarizer.eval() | |
| translator.to(device) | |
| summarizer.to(device) | |
| # def generate_multiple_choice_question( | |
| # context | |
| # ): | |
| # num_questions = 1 | |
| # question_item = question_generation_sampling( | |
| # g1_model, g1_tokenizer, | |
| # g2_model, g2_tokenizer, | |
| # context, num_questions, device | |
| # )[0] | |
| # question = question_item['question'] | |
| # options = question_item['options'] | |
| # options[0] = f"{options[0]} [ANSWER]" | |
| # random.shuffle(options) | |
| # output_string = f"Question: {question}\n[A] {options[0]}\n[B] {options[1]}\n[C] {options[2]}\n[D] {options[3]}" | |
| # return output_string | |
| # | |
| # demo = gr.Interface( | |
| # fn=generate_multiple_choice_question, | |
| # inputs=gr.Textbox(lines=8, placeholder="Context Here..."), | |
| # outputs=gr.Textbox(lines=5, placeholder="Question: \n[A] \n[B] \n[C] \n[D] "), | |
| # title="Multiple-choice Question Generator", | |
| # description="Provide some context (e.g. news article or any passage) in the context box and click **Submit**. The models currently support English only. This demo is a part of MQAG - https://github.com/potsawee/mqag0.", | |
| # allow_flagging='never' | |
| # ) | |
| def generate_output( | |
| task, | |
| text, | |
| ): | |
| inputs = tokenizer( | |
| [text], | |
| padding="longest", | |
| max_length=1024, | |
| truncation=True, | |
| return_tensors="pt", | |
| ).to(device) | |
| if task == 'Translation': | |
| outputs = translator.generate( | |
| **inputs, | |
| max_new_tokens=256, | |
| ) | |
| elif task == 'Summarization': | |
| outputs = summarizer.generate( | |
| **inputs, | |
| max_new_tokens=256, | |
| ) | |
| else: | |
| raise ValueError("task undefined!") | |
| gen_text = tokenizer.decode(outputs[0], skip_special_tokens=True) | |
| return gen_text | |
| TASKS = ["Translation", "Summarization"] | |
| demo = gr.Interface( | |
| fn=generate_output, | |
| inputs=[ | |
| gr.components.Radio(label="Task", choices=TASKS, value="Translation"), | |
| gr.components.Textbox(label="Text (in English)", lines=10), | |
| ], | |
| outputs=gr.Textbox(label="Text (in Thai)", lines=4), | |
| # examples=[["Building a translation demo with Gradio is so easy!", "eng_Latn", "spa_Latn"]], | |
| cache_examples=False, | |
| title="English🇬🇧 to Thai🇹🇭 | Translation or Summarization", | |
| description="Provide some text (in English) & select one of the tasks (Translation or Summarization). Note that currently the model only supports text up to 1024 tokens. The base architecture is mt5-large with the embeddings filtered to only English and Thai tokens and fine-tuned to XSum (Eng2Thai) Dataset (https://huggingface.co/datasets/potsawee/xsum_eng2thai).", | |
| allow_flagging='never' | |
| ) | |
| demo.launch() | |