Spaces:
Sleeping
Sleeping
| import os, sys | |
| # from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, MBartForConditionalGeneration | |
| # import torch | |
| import gradio as gr | |
| import requests | |
| import json | |
| # from huggingface_hub import login | |
| class LTRC_Translation_API(): | |
| def __init__(self, url = 'https://ssmt.iiit.ac.in/onemt', src_lang = 'en', tgt_lang = 'te'): | |
| self.lang_map = {'te': 'tel', 'en': 'eng', 'ta': 'tam', 'ml': 'mal', 'mr': 'mar', 'kn': 'kan', 'hi': 'hin'} | |
| self.url = url | |
| self.headers = { | |
| 'Content-Type': 'application/json', | |
| 'Accept': 'application/json' | |
| } | |
| tgt_lang = self.lang_map.get(tgt_lang, 'te') | |
| src_lang = self.lang_map.get(src_lang, 'en') | |
| self.src_lang = src_lang | |
| self.tgt_lang = tgt_lang | |
| def translate(self, text): | |
| try: | |
| data = {'text': text, 'source_language': self.src_lang, 'target_language': self.tgt_lang} | |
| response = requests.post(self.url, headers = self.headers, json = data) | |
| translated_text = json.loads(response.text).get('data', '') | |
| return translated_text | |
| except Exception as e: | |
| print("Exception: ", e) | |
| return '' | |
| # class Headline_Generation(): | |
| # def __init__(self, model_name = "lokeshmadasu42/sample"): | |
| # self.model_name = model_name | |
| # self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") | |
| # self.tokenizer = AutoTokenizer.from_pretrained(model_name, do_lower_case=False, use_fast=False, keep_accents=True) | |
| # self.model = AutoModelForSeq2SeqLM.from_pretrained(model_name) | |
| # self.model.to(self.device) | |
| # self.model.eval() | |
| # self.bos_id = self.tokenizer._convert_token_to_id_with_added_voc("<s>") | |
| # self.eos_id = self.tokenizer._convert_token_to_id_with_added_voc("</s>") | |
| # self.pad_id = self.tokenizer._convert_token_to_id_with_added_voc("<pad>") | |
| # self.lang_map = {'as': '<2as>', 'bn': '<2bn>', 'en': '<2en>', 'gu': '<2gu>', 'hi': '<2hi>', 'kn': '<2kn>', 'ml': '<2ml>', 'mr': '<2mr>', 'or': '<2or>', 'pa': '<2pa>', 'ta': '<2ta>', 'te': '<2te>'} | |
| # print("Headline Generation model loaded...!") | |
| # def get_headline(self, text, lang_id): | |
| # inp = self.tokenizer(text, add_special_tokens=False, return_tensors="pt", padding=True).to(self.device) | |
| # inp = inp['input_ids'] | |
| # lang_code = self.lang_map.get(lang_id, '') | |
| # text = text + "</s> " + lang_code | |
| # # print("Text: ", text) | |
| # model_output = self.model.generate( | |
| # inp, | |
| # use_cache=True, | |
| # num_beams=5, | |
| # max_length=32, | |
| # min_length=1, | |
| # early_stopping=True, | |
| # pad_token_id = self.pad_id, | |
| # bos_token_id = self.bos_id, | |
| # eos_token_id = self.eos_id, | |
| # decoder_start_token_id = self.tokenizer._convert_token_to_id_with_added_voc(lang_code) | |
| # ) | |
| # decoded_output = self.tokenizer.decode( | |
| # model_output[0], | |
| # skip_special_tokens=True, | |
| # clean_up_tokenization_spaces=False | |
| # ) | |
| # return decoded_output | |
| # class Summarization(): | |
| # def __init__(self, model_name = "ashokurlana/mBART-TeSum"): | |
| # self.model_name = model_name | |
| # self.device = torch.device("cuda:1" if torch.cuda.is_available() else "cpu") | |
| # self.tokenizer = AutoTokenizer.from_pretrained(model_name) | |
| # self.model = AutoModelForSeq2SeqLM.from_pretrained(model_name) | |
| # self.model.to(self.device) | |
| # self.model.eval() | |
| # self.lang_map = {'te': 'te_IN', 'en': 'en_XX'} | |
| # print("Summarization model loaded...!") | |
| # def get_summary(self, text, lang_id): | |
| # inp = self.tokenizer([text], add_special_tokens=False, return_tensors="pt", max_length = 1024).to(self.device) | |
| # inp = inp['input_ids'] | |
| # lang_code = self.lang_map.get(lang_id, '') | |
| # model_output = self.model.generate( | |
| # inp, | |
| # use_cache=True, | |
| # num_beams=5, | |
| # max_length=256, | |
| # early_stopping=True | |
| # ) | |
| # decoded_output = [self.tokenizer.decode( | |
| # summ_id, | |
| # skip_special_tokens=True, | |
| # clean_up_tokenization_spaces=False | |
| # ) for summ_id in model_output] | |
| # return " ".join(decoded_output) | |
| def get_prediction(text, src_lang_id, tgt_lang_id, translate = False): | |
| # if len(sys.argv)<3: | |
| # print("Usage: python app.py <text_file_path> <lang_id>") | |
| # print("Text file should contain the article news") | |
| # exit() | |
| # txt_path = sys.argv[1] | |
| # lang_id = sys.argv[2] | |
| # if not os.path.exists(txt_path): | |
| # print("Path: {} do not exists".format(txt_path)) | |
| # exit() | |
| # text = '' | |
| # with open(txt_path, 'r', encoding='utf-8') as fp: | |
| # text = fp.read().strip() | |
| ### Login to huggingface token | |
| # access_token = "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" | |
| # login(token = access_token) | |
| # headline_generator = Headline_Generation() | |
| # summarizer = Summarization() | |
| # if translate == True: | |
| # translator = LTRC_Translation_API(tgt_lang = lang_id) | |
| # text = translator.translate(text) | |
| # headline = headline_generator.get_headline(text, lang_id) | |
| # summary = summarizer.get_summary(text, lang_id) | |
| # print("Article: ", text) | |
| # print("Summary: ", summary) | |
| # print("Headline: ", headline) | |
| # return "Headline: " + headline + "\nSummary: " + summary | |
| # return [text, summary, headline] | |
| translator = LTRC_Translation_API(src_lang = src_lang_id, tgt_lang = tgt_lang_id) | |
| text = translator.translate(text) | |
| return text | |
| interface = gr.Interface( | |
| get_prediction, | |
| inputs=[ | |
| gr.Textbox(lines = 8, label = "Source Text", info = "Provide the news article text here"), | |
| # gr.Textbox(lines = 8, label = "News Article Text", info = "Provide the news article text here. Check the `Translate` if the source language is english."), | |
| # gr.Dropdown( | |
| # ['as', 'bn', 'en', 'gu', 'hi', 'kn', 'ml', 'mr', 'or', 'pa', 'ta', 'te'], label="Source Language code", info="select the source language code" | |
| # ), | |
| # gr.Dropdown( | |
| # ['as', 'bn', 'en', 'gu', 'hi', 'kn', 'ml', 'mr', 'or', 'pa', 'ta', 'te'], label="Target Language code", info="select the target language code" | |
| # ), | |
| gr.Dropdown( | |
| ['en', 'hi', 'kn', 'ml', 'mr', 'ta', 'te'], label="Source Language code", info="select the source language code" | |
| ), | |
| gr.Dropdown( | |
| ['en', 'hi', 'kn', 'ml', 'mr', 'ta', 'te'], label="Target Language code", info="select the target language code" | |
| ), | |
| # gr.Checkbox(label="Translate", info="Is translation required?") | |
| ], | |
| outputs=[ | |
| gr.Textbox(lines = 8, label = "Translation", info = "Translated text"), | |
| # gr.Textbox(lines = 8, label = "Source Article Text", info = "Source article text (if `Translate` is enabled then the source will be translated to target language)"), | |
| # gr.Textbox(lines = 4, label = "Summary", info = "Summary of the given article (translated if `Translate` is enabled)"), | |
| # gr.Textbox(lines = 2, label = "Headline", info = "Generated headline of the given article (translated if `Translate` is enabled)") | |
| ], | |
| title = "Indic Translation Demo" | |
| ) | |
| interface.launch(share=True) |