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()