import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

# Load Anki-2.5 model
model_name = "anktechsol/anki-2.5"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float32,
    device_map="auto"
)

def respond(message, history, system_message, max_tokens, temperature, top_p):
    """
    Market Intelligence Assistant for Indian Businesses powered by Anki-2.5
    Provides insights on Indian market trends, business strategies, and more
    """
    # Build prompt manually
    prompt = system_message + "\n\n"
    
    # Add history
    for msg in history:
        if isinstance(msg, dict):
            role = msg.get("role", "")
            content = msg.get("content", "")
            if role == "user":
                prompt += f"User: {content}\n"
            elif role == "assistant":
                prompt += f"Assistant: {content}\n"
        elif len(msg) == 2:
            prompt += f"User: {msg[0]}\n"
            prompt += f"Assistant: {msg[1]}\n"
    
    # Add current message
    prompt += f"User: {message}\nAssistant:"
    
    # Generate response
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    
    with torch.no_grad():
        outputs = model.generate(
            inputs["input_ids"],
            max_new_tokens=max_tokens,
            temperature=temperature,
            top_p=top_p,
            do_sample=True,
            pad_token_id=tokenizer.eos_token_id
        )
    
    response = tokenizer.decode(outputs[0][inputs["input_ids"].shape[-1]:], skip_special_tokens=True)
    return response

# Create ChatInterface focused on business intelligence
chatbot = gr.ChatInterface(
    respond,
    title="📈 Indian Market Intelligence - Powered by Anki-2.5",
    description="Get market insights, business strategies, and intelligence for the Indian market. Ask about trends, regulations, consumer behavior, and more!",
    examples=[
        ["भारत में ई-कॉमर्स बिजनेस शुरू करने के लिए क्या चाहिए?"],
        ["What are the key consumer trends in tier 2 Indian cities?"],
        ["Explain GST compliance for small businesses in India"],
        ["ভারতীয় বাজারে ডিজিটাল মার্কেটিং এর গুরুত্ব"],
        ["How can I expand my retail business to rural India?"],
    ],
    additional_inputs=[
        gr.Textbox(
            value="You are a market intelligence assistant specializing in the Indian market. You provide insights on business trends, consumer behavior, regulations like GST, market entry strategies, and digital transformation. You understand the Indian business landscape including urban, tier-2, and rural markets. Respond in Hindi or English as needed.",
            label="System Message"
        ),
        gr.Slider(minimum=50, maximum=1024, value=512, step=1, label="Max Tokens"),
        gr.Slider(minimum=0.1, maximum=1.5, value=0.7, step=0.1, label="Temperature"),
        gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p"),
    ],
    theme=gr.themes.Soft(),
)

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