KnowTheText / app.py
SumanthKarnati's picture
Update app.py
96e47fa
# Import necessary libraries
import streamlit as st
from transformers import AutoModelForSequenceClassification, AutoTokenizer, AutoModelForSeq2SeqLM
import torch
# Define function to load models
@st.cache_data(allow_output_mutation=True)
def load_models():
classification_model_name = 'distilbert-base-uncased-finetuned-sst-2-english'
classification_model = AutoModelForSequenceClassification.from_pretrained(classification_model_name)
classification_tokenizer = AutoTokenizer.from_pretrained(classification_model_name, model_max_length=512)
summarization_model_name = 't5-base'
summarization_model = AutoModelForSeq2SeqLM.from_pretrained(summarization_model_name)
summarization_tokenizer = AutoTokenizer.from_pretrained(summarization_model_name, model_max_length=512)
return classification_model, classification_tokenizer, summarization_model, summarization_tokenizer
classification_model, classification_tokenizer, summarization_model, summarization_tokenizer = load_models()
# Title of the app
st.title('Text Classification and Summarization with Hugging Face')
# Take user input
text = st.text_area("Enter text:", "")
submit_button = st.button("Analyze Text")
# Predict function for sentiment analysis
def predict_sentiment(text):
tokenized_text = classification_tokenizer.tokenize(text)
results = []
# Break text into chunks of max_model_length tokens
for i in range(0, len(tokenized_text), classification_tokenizer.model_max_length):
chunk = tokenized_text[i:i+classification_tokenizer.model_max_length]
chunk = classification_tokenizer.convert_tokens_to_string(chunk)
inputs = classification_tokenizer(chunk, return_tensors="pt", truncation=True, padding='max_length')
outputs = classification_model(**inputs)
probs = torch.nn.functional.softmax(outputs[0], dim=-1)
results.append(probs.detach().numpy())
return results
# Predict function for text summarization
def summarize_text(text):
tokenized_text = summarization_tokenizer.tokenize(text)
summaries = []
# Break text into chunks of max_model_length tokens
for i in range(0, len(tokenized_text), summarization_tokenizer.model_max_length):
chunk = tokenized_text[i:i+summarization_tokenizer.model_max_length]
chunk = summarization_tokenizer.convert_tokens_to_string(chunk)
inputs = summarization_tokenizer.encode("summarize: " + chunk, return_tensors="pt", truncation=True, padding='max_length')
outputs = summarization_model.generate(inputs, max_length=150, min_length=40, length_penalty=2.0, num_beams=4, early_stopping=True)
summary = summarization_tokenizer.decode(outputs[0]).replace('<pad>', '').replace('</s>', '')
summaries.append(summary)
return summaries
if submit_button:
if text:
with st.spinner("Analyzing..."):
# Sentiment analysis
results = predict_sentiment(text)
for i, probs in enumerate(results):
st.markdown(f"**Result {i+1}:**")
st.markdown(f"**Positive sentiment:** `{probs[0][1]:.2f}`")
st.markdown(f"**Negative sentiment:** `{probs[0][0]:.2f}`")
# Text summarization
summaries = summarize_text(text)
for i, summary in enumerate(summaries):
st.markdown(f"**Summary {i+1}:** `{summary}`")
else:
st.warning("Please enter text to analyze.")