Spaces:
Sleeping
Sleeping
import gradio as gr | |
from transformers import pipeline | |
from loguru import logger | |
# from pydantic import BaseModel | |
# RU_SUMMARY_MODEL = "IlyaGusev/rubart-large-sum" | |
RU_SUMMARY_MODEL = "IlyaGusev/mbart_ru_sum_gazeta" | |
# RU_SENTIMENT_MODEL = "IlyaGusev/rubart-large-sentiment" | |
RU_SENTIMENT_MODEL = "seara/rubert-tiny2-russian-sentiment" | |
EN_SUMMARY_MODEL = "sshleifer/distilbart-cnn-12-6" | |
EN_SENTIMENT_MODEL = "ProsusAI/finbert" | |
class Summarizer(): | |
ru_summary_pipe: pipeline | |
ru_sentiment_pipe: pipeline | |
def __init__(self) -> None: | |
self.ru_summary_pipe = pipeline("summarization", model=RU_SUMMARY_MODEL, max_length=100, truncation=True) | |
self.ru_sentiment_pipe = pipeline("sentiment-analysis", model=RU_SENTIMENT_MODEL) | |
def summarize(self, text: str) -> str: | |
result = {} | |
response_summary = self.ru_summary_pipe(text) | |
logger.info(response_summary) | |
result["summary"] = response_summary[0]["summary_text"] | |
response_sentiment = self.ru_sentiment_pipe(text) | |
logger.info(response_sentiment) | |
result["sentiment"] = response_sentiment[0]["label"] | |
return f"Summary: {result['summary']}\n Sentiment:{result['sentiment']}" | |
pipe = Summarizer() | |
demo = gr.Interface( | |
fn=pipe.summarize, | |
inputs=gr.Textbox(lines=5, placeholder="Write your text here..."), | |
outputs=gr.Textbox(lines=5, placeholder="Summary and Sentiment would be here..."), | |
) | |
if __name__ == "__main__": | |
demo.launch(show_api=False) |