NLP-Model / pipeline.py
iamomtiwari's picture
Create pipeline.py
786d2c0 verified
from transformers import T5ForConditionalGeneration, T5Tokenizer
from googletrans import Translator
class GOTSummarizer:
def __init__(self):
self.model = T5ForConditionalGeneration.from_pretrained("yourusername/got_summarizer")
self.tokenizer = T5Tokenizer.from_pretrained("yourusername/got_summarizer")
self.translator = Translator()
def summarize(self, text, max_length=150):
inputs = self.tokenizer("summarize: " + text, return_tensors="pt", truncation=True)
outputs = self.model.generate(**inputs, max_length=max_length)
return self.tokenizer.decode(outputs[0], skip_special_tokens=True)
def translate(self, text, lang='hi'):
return self.translator.translate(text, dest=lang).text