Nitin00043's picture
Update app.py
9024f72 verified
# from transformers import pipeline, AutoModelForSequenceClassification, AutoTokenizer
# import gradio as gr
# import torch
# from concurrent.futures import ThreadPoolExecutor
# from threading import Lock
# # Global cache and lock for thread-safety
# CACHE_SIZE = 100
# prediction_cache = {}
# cache_lock = Lock()
# # Mapping for sentiment labels from cardiffnlp/twitter-roberta-base-sentiment
# SENTIMENT_LABEL_MAPPING = {
# "LABEL_0": "negative",
# "LABEL_1": "neutral",
# "LABEL_2": "positive"
# }
# def load_model(model_name):
# """
# Loads the model with 8-bit quantization if a GPU is available;
# otherwise, loads the full model.
# """
# if torch.cuda.is_available():
# model = AutoModelForSequenceClassification.from_pretrained(
# model_name, load_in_8bit=True, device_map="auto"
# )
# tokenizer = AutoTokenizer.from_pretrained(model_name)
# device = 0 # GPU index
# else:
# model = AutoModelForSequenceClassification.from_pretrained(model_name)
# tokenizer = AutoTokenizer.from_pretrained(model_name)
# device = -1
# return pipeline("text-classification", model=model, tokenizer=tokenizer, device=device)
# # Load both models concurrently at startup.
# with ThreadPoolExecutor() as executor:
# sentiment_future = executor.submit(load_model, "cardiffnlp/twitter-roberta-base-sentiment")
# emotion_future = executor.submit(load_model, "bhadresh-savani/bert-base-uncased-emotion")
# sentiment_pipeline = sentiment_future.result()
# emotion_pipeline = emotion_future.result()
# def analyze_text(text):
# # Check cache first (thread-safe)
# with cache_lock:
# if text in prediction_cache:
# return prediction_cache[text]
# try:
# # Run both model inferences in parallel.
# with ThreadPoolExecutor() as executor:
# future_sentiment = executor.submit(sentiment_pipeline, text)
# future_emotion = executor.submit(emotion_pipeline, text)
# sentiment_result = future_sentiment.result()[0]
# emotion_result = future_emotion.result()[0]
# # Remap the sentiment label to a human-readable format if available.
# raw_sentiment_label = sentiment_result.get("label", "")
# sentiment_label = SENTIMENT_LABEL_MAPPING.get(raw_sentiment_label, raw_sentiment_label)
# # Format the output with rounded scores.
# result = {
# "Sentiment": {sentiment_label: round(sentiment_result['score'], 4)},
# "Emotion": {emotion_result['label']: round(emotion_result['score'], 4)}
# }
# except Exception as e:
# result = {"error": str(e)}
# # Update the cache in a thread-safe manner.
# with cache_lock:
# if len(prediction_cache) >= CACHE_SIZE:
# prediction_cache.pop(next(iter(prediction_cache)))
# prediction_cache[text] = result
# return result
# # Define the Gradio interface.
# demo = gr.Interface(
# fn=analyze_text,
# inputs=gr.Textbox(placeholder="Enter your text here...", label="Input Text"),
# outputs=gr.JSON(label="Analysis Results"),
# title="πŸš€ Fast Sentiment & Emotion Analysis",
# description="Optimized application that remaps sentiment labels and uses 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 the models with a sample input.
# _ = analyze_text("Warming up models...")
# if __name__ == "__main__":
# # Bind to all interfaces for Hugging Face Spaces.
# demo.launch(server_name="0.0.0.0", server_port=7860)