Update app.py

app.py

@@ -1,10 +1,10 @@
 import os, sys
-from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, MBartForConditionalGeneration
-import torch
+# from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, MBartForConditionalGeneration
+# import torch
 import gradio as gr
 import requests
 import json
-from huggingface_hub import login
+# from huggingface_hub import login
 
 
 class LTRC_Translation_API():
@@ -37,94 +37,94 @@ class LTRC_Translation_API():
             return ''
             
 
-class Headline_Generation():
-	def __init__(self, model_name = "lokeshmadasu42/sample"):
-		self.model_name = model_name
+# 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.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.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>'}
+# 		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...!")
+# 		print("Headline Generation model loaded...!")
 
 
-	def get_headline(self, text, lang_id):
+# 	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']
+# 		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, '')
+# 		lang_code = self.lang_map.get(lang_id, '')
 
-		text = text + "</s> " + lang_code
-		# print("Text: ", text)
+# 		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)
-		)
+# 		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
-		)
+# 		decoded_output = self.tokenizer.decode(
+# 			model_output[0],
+# 			skip_special_tokens=True,
+# 			clean_up_tokenization_spaces=False
+# 		)
 
-		return decoded_output
+# 		return decoded_output
 
 
-class Summarization():
-	def __init__(self, model_name = "ashokurlana/mBART-TeSum"):
-		self.model_name = model_name
+# 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.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'}
+# 		self.lang_map = {'te': 'te_IN', 'en': 'en_XX'}
 
-		print("Summarization model loaded...!")
+# 		print("Summarization model loaded...!")
 
 
-	def get_summary(self, text, lang_id):
+# 	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']
+# 		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, '')
+# 		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
-		)
+# 		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]
+# 		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)
+# 		return " ".join(decoded_output)
 
 
 def get_prediction(text, lang_id, translate = False):