NLP-Model / app.py
iamomtiwari's picture
Update app.py
71f7ef1 verified
import subprocess
import sys
from transformers import T5ForConditionalGeneration, T5Tokenizer
import gradio as gr
# Install required packages
def install(package):
subprocess.check_call([sys.executable, "-m", "pip", "install", package])
install('transformers')
install('sentencepiece')
install('torch')
install('gradio')
# Translation setup
try:
install('translate')
from translate import Translator
translation_available = True
except:
print("Translation package not available - continuing without translation support")
translation_available = False
class GOTSummarizer:
def __init__(self):
self.tokenizer = T5Tokenizer.from_pretrained("t5-small")
self.model = T5ForConditionalGeneration.from_pretrained("t5-small")
self.translation_available = translation_available
def summarize(self, text, max_length=150):
inputs = self.tokenizer("summarize: " + text,
return_tensors="pt",
truncation=True,
max_length=512)
outputs = self.model.generate(
**inputs,
max_length=max_length,
num_beams=4,
early_stopping=True
)
return self.tokenizer.decode(outputs[0], skip_special_tokens=True)
def translate(self, text, lang='hi'):
if not self.translation_available:
return "[Translation unavailable] " + text
try:
translator = Translator(to_lang=lang) # Initialize with target language
return translator.translate(text)
except Exception as e:
print(f"Translation error: {e}")
return text
# Initialize summarizer
summarizer = GOTSummarizer()
def process(text, lang='en'):
if not text.strip():
return "Please enter text to summarize"
summary = summarizer.summarize(text)
return summary if lang == 'en' else summarizer.translate(summary, lang)
# Create Gradio interface
interface = gr.Interface(
fn=process,
inputs=[
gr.Textbox(label="Game of Thrones Text", lines=10,
placeholder="Paste book chapter text here..."),
gr.Dropdown(
label="Language",
choices=['en', 'hi', 'pa', 'ta', 'bn'],
value='en'
)
],
outputs=gr.Textbox(label="Summary"),
title="ASOIAF Chapter Summarizer",
description="Summarizes Game of Thrones chapters with optional translation"
)
# Launch the interface
if __name__ == "__main__":
interface.launch()