import gradio as gr
from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification
from groq import Groq
import os

# Set Groq API key
os.environ["GROQ_API_KEY"] = "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"

# Text Classification Setup
classifier = pipeline("zero-shot-classification",
                      model="facebook/bart-large-mnli")

# Chatbot Setup
client = Groq()
system_prompt = """You are an advanced AI assistant with deep contextual understanding. 
Maintain natural conversation while demonstrating:
1. Complex sentence comprehension
2. Contextual awareness across multiple turns
3. Emotional intelligence
4. Domain-specific knowledge adaptation"""

def classify_text(text, labels):
    labels = [label.strip() for label in labels.split(",")]
    results = classifier(text, labels, multi_label=False)
    return {label: score for label, score in zip(results["labels"], results["scores"])}

def groq_chat(user_input, history):
    conversation = [{"role": "system", "content": system_prompt}]
    
    for user, assistant in history:
        conversation.extend([
            {"role": "user", "content": user},
            {"role": "assistant", "content": assistant}
        ])
    
    conversation.append({"role": "user", "content": user_input})
    
    response = client.chat.completions.create(
        model="llama3-70b-8192",
        messages=conversation,
        temperature=0.7,
        max_tokens=512,
        top_p=1
    )
    
    return response.choices[0].message.content

# Gradio Interface
with gr.Blocks() as app:
    gr.Markdown("# Advanced LLM Application")
    
    with gr.Tab("Text Classification"):
        with gr.Row():
            with gr.Column():
                text_input = gr.Textbox(label="Input Text")
                labels_input = gr.Textbox(label="Categories (comma-separated)", 
                                        value="positive, negative, neutral")
                classify_btn = gr.Button("Classify")
            results_output = gr.Label(label="Classification Results")
        
        classify_btn.click(
            fn=classify_text,
            inputs=[text_input, labels_input],
            outputs=results_output
        )
    
    with gr.Tab("Chatbot"):
        chatbot = gr.Chatbot(height=400)
        msg = gr.Textbox(label="Your Message")
        clear = gr.Button("Clear")
        
        def respond(message, chat_history):
            bot_message = groq_chat(message, chat_history)
            chat_history.append((message, bot_message))
            return "", chat_history
        
        msg.submit(respond, [msg, chatbot], [msg, chatbot])
        clear.click(lambda: None, None, chatbot, queue=False)

app.launch()