from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
import torch
import gradio as gr
from peft import PeftModel, PeftConfig
import spaces

# Use the GPU if available
device = 0 if torch.cuda.is_available() else -1

def load_model():
    base_model_name = "Qwen/Qwen2.5-1.5B-Instruct"
    tokenizer = AutoTokenizer.from_pretrained(base_model_name)
    base_model = AutoModelForCausalLM.from_pretrained(base_model_name)
    
    peft_model = PeftModel.from_pretrained(
        base_model,
        "ombhojane/smile-small",
    )
    
    return pipeline(
        "text-generation",
        model=peft_model,
        tokenizer=tokenizer,
        device=device
    )

pipe = load_model()

@spaces.GPU
def generate_response(message, max_length, temperature):
    if not message:
        return "Please enter a message to generate a response."
    
    messages = [
        {"role": "user", "content": message}
    ]
    
    generated_text = pipe(
        messages, 
        max_new_tokens=max_length,
        num_return_sequences=1,
        temperature=temperature,
        do_sample=True
    )
    
    response = generated_text[0]['generated_text']
    if isinstance(response, list):
        for msg in response:
            if msg.get('role') == 'assistant':
                return msg.get('content', '')
    return response

# Custom CSS for better styling
custom_css = """
.gradio-container {
    font-family: 'Arial', sans-serif;
}
.main-title {
    text-align: center;
    color: #2C3E50;
    margin-bottom: 1em;
}
.description {
    text-align: justify;
    margin-bottom: 2em;
    color: #34495E;
}
"""

with gr.Blocks(css=custom_css) as demo:
    gr.Markdown(
        """
        # 🌟 S.M.I.L.E for India
        ### Smart Marketing Intelligence & Local Engagement
        """
    )
    
    with gr.Row():
        with gr.Column(scale=2):
            gr.Markdown(
                """
                ### About
                S.M.I.L.E is an AI assistant trained to handle Indian customer queries with cultural context awareness, 
                beneficial for Indian CRMs. It understands local nuances, cultural preferences, and consumer behaviors 
                specific to the Indian market.
                """
            )
    
    with gr.Row():
        with gr.Column(scale=2):
            input_text = gr.Textbox(
                lines=4,
                placeholder="Enter your query about Indian consumer behavior, cultural aspects, or market trends...",
                label="Your Query"
            )
            
            with gr.Row():
                max_length = gr.Slider(
                    minimum=50,
                    maximum=500,
                    value=200,
                    step=10,
                    label="Maximum Response Length",
                    info="Adjust the length of the generated response"
                )
                temperature = gr.Slider(
                    minimum=0.1,
                    maximum=1.0,
                    value=0.7,
                    step=0.1,
                    label="Temperature",
                    info="Higher values make the output more creative, lower values make it more focused"
                )
            
            submit_btn = gr.Button("Generate Response", variant="primary")
            
        with gr.Column(scale=2):
            output_text = gr.Textbox(
                lines=12,
                label="AI Response",
                show_copy_button=True
            )
            
    with gr.Accordion("Sample Queries", open=False):
        gr.Markdown(
            """
            - What are the key factors influencing Indian consumer buying behavior?
            - How do cultural values affect marketing strategies in India?
            - What are the popular festival shopping trends in India?
            - Explain the role of family decision-making in Indian purchases
            """
        )
    
    submit_btn.click(
        generate_response,
        inputs=[input_text, max_length, temperature],
        outputs=output_text
    )

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