import gradio as gr
from transformers import pipeline
import torch

# Load the zero-shot classification model
try:
    model_name = "MoritzLaurer/DeBERTa-v3-large-mnli-fever-anli-ling-wanli"
    classifier = pipeline("zero-shot-classification", 
                        model=model_name, 
                        device=0 if torch.cuda.is_available() else -1)
except Exception as e:
    print(f"Error loading main model: {e}")
    # Fallback to a lighter model if the first one fails
    model_name = "facebook/bart-large-mnli"
    classifier = pipeline("zero-shot-classification", model=model_name)

def classify_product(ad_text):
    if not ad_text.strip():
        return "Please enter some ad text."
    
    try:
        # Category classification
        category_result = classifier(
            ad_text,
            candidate_labels=[
                "Software", "Electronics", "Clothing", "Food & Beverage", 
                "Healthcare", "Financial Services", "Entertainment",
                "Home & Garden", "Automotive", "Education"
            ],
            hypothesis_template="This is an advertisement for a product in the {} category",
            multi_label=False
        )
        
        # Product type classification
        product_result = classifier(
            ad_text,
            candidate_labels=[
                "software application", "mobile app", "subscription service",
                "physical product", "digital product", "professional service",
                "consumer device", "platform", "tool"
            ],
            hypothesis_template="This is specifically a {}",
            multi_label=False
        )
        
        # Format output string
        output = f"""
📊 Analysis Results:

🏷️ Category: {category_result['labels'][0]}
   Confidence: {category_result['scores'][0]:.2%}

📦 Product Type: {product_result['labels'][0]}
   Confidence: {product_result['scores'][0]:.2%}
"""
        
        # Additional product details from text
        if any(brand_keyword in ad_text.lower() for brand_keyword in ['by', 'from', 'introducing', 'new']):
            product_name_result = classifier(
                ad_text,
                candidate_labels=["contains brand name", "does not contain brand name"],
                hypothesis_template="This text {}",
                multi_label=False
            )
            if product_name_result['labels'][0] == "contains brand name":
                output += "\n🏢 Brand Mention: Likely contains a brand name"
        
        return output
    
    except Exception as e:
        return f"An error occurred: {str(e)}\nPlease try with different text or contact support."

# Create Gradio interface
demo = gr.Interface(
    fn=classify_product,
    inputs=gr.Textbox(
        lines=5,
        placeholder="Paste your ad text here (max 100 words)...",
        label="Advertisement Text"
    ),
    outputs=gr.Textbox(label="Analysis Results"),
    title="AI Powered Product Identifier from Ad Text",
    description="Paste your marketing ad text to identify the product category and type. Maximum 100 words.",
    examples=[
        ["Experience seamless productivity with our new CloudWork Pro subscription. This AI-powered workspace solution helps remote teams collaborate better with smart document sharing, real-time editing, and integrated chat features. Starting at $29/month."],
        ["Introducing the new iPhone 15 Pro with revolutionary A17 Pro chip. Capture stunning photos with our advanced 48MP camera system. Available in titanium finish with all-day battery life. Pre-order now at Apple stores nationwide."],
    ],
    theme=gr.themes.Soft()
)

if __name__ == "__main__":
    demo.launch()