Nitin00043's picture
Update app.py
adc05de verified
raw
history blame
2.57 kB
from transformers import pipeline, AutoModelForSequenceClassification, AutoTokenizer
import gradio as gr
import torch
from concurrent.futures import ThreadPoolExecutor
# Load models with quantization (8-bit) for faster inference
def load_quantized_model(model_name):
model = AutoModelForSequenceClassification.from_pretrained(model_name, load_in_8bit=True)
tokenizer = AutoTokenizer.from_pretrained(model_name)
return pipeline("text-classification", model=model, tokenizer=tokenizer, device=0 if torch.cuda.is_available() else -1)
# Load models in parallel during startup
with ThreadPoolExecutor() as executor:
sentiment_future = executor.submit(load_quantized_model, "cardiffnlp/twitter-roberta-base-sentiment")
emotion_future = executor.submit(load_quantized_model, "bhadresh-savani/bert-base-uncased-emotion")
sentiment_pipeline = sentiment_future.result()
emotion_pipeline = emotion_future.result()
# Cache recent predictions to avoid recomputation
CACHE_SIZE = 100
prediction_cache = {}
def analyze_text(text):
# Check cache first
if text in prediction_cache:
return prediction_cache[text]
# Parallel model execution
with ThreadPoolExecutor() as executor:
sentiment_future = executor.submit(sentiment_pipeline, text)
emotion_future = executor.submit(emotion_pipeline, text)
sentiment_result = sentiment_future.result()[0]
emotion_result = emotion_future.result()[0]
# Format response
result = {
"Sentiment": {sentiment_result['label']: round(sentiment_result['score'], 4)},
"Emotion": {emotion_result['label']: round(emotion_result['score'], 4)}
}
# Update cache
if len(prediction_cache) >= CACHE_SIZE:
prediction_cache.pop(next(iter(prediction_cache)))
prediction_cache[text] = result
return result
# Optimized Gradio interface with batch processing
demo = gr.Interface(
fn=analyze_text,
inputs=gr.Textbox(placeholder="Enter your text here...", label="Input Text"),
outputs=gr.Label(label="Analysis Results"),
title="πŸš€ Fast Sentiment & Emotion Analysis",
description="Optimized version using quantized models and parallel processing",
examples=[
["I'm thrilled to start this new adventure!"],
["This situation is making me really frustrated."],
["I feel so heartbroken and lost."]
],
theme="soft",
allow_flagging="never"
)
# Warm up models with sample input
analyze_text("Warming up models...")
if __name__ == "__main__":
demo.launch()