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

# Model configuration
MODEL_ID = "mx-llms/Lychee-GPT-9B"
MAX_NEW_TOKENS = 256
TEMPERATURE = 0.7
TOP_P = 0.9

# Load model and tokenizer
print("Loading model... (এটা প্রথমবার একটু সময় লাগবে)")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)

# CPU optimization - 8-bit quantization
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    device_map="cpu",  # CPU তে run করবে
    load_in_8bit=False,  # CPU তে 8bit সাপোর্ট নেই, তাই False
    torch_dtype=torch.float32,  # CPU তে float32 ব্যবহার করতে হয়
    trust_remote_code=True,
)

# Set to eval mode
model.eval()

def generate_response(user_message, temperature=0.7, top_p=0.9, max_tokens=256):
    """Generate response from Lychee-GPT"""
    try:
        # Prepare input
        inputs = tokenizer(user_message, return_tensors="pt")
        
        # Generate
        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,
                eos_token_id=tokenizer.eos_token_id,
            )
        
        # Decode
        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
        
        # Remove the input from response
        if user_message in response:
            response = response.replace(user_message, "", 1).strip()
        
        return response
    
    except Exception as e:
        return f"Error: {str(e)}"

# Gradio Interface
with gr.Blocks(title="Lychee-GPT-9B") as demo:
    gr.Markdown("""
    # 🎉 Lychee-GPT-9B
    ### আপনার নিজস্ব LLM Model!
    
    > ⚠️ **নোট:** CPU তে চলছে তাই response 30-90 সেকেন্ড লাগতে পারে।
    """)
    
    with gr.Row():
        with gr.Column():
            user_input = gr.Textbox(
                label="আপনার প্রশ্ন/বার্তা",
                placeholder="কিছু লিখুন...",
                lines=3
            )
            
            with gr.Row():
                temp_slider = gr.Slider(
                    label="Temperature",
                    minimum=0.0,
                    maximum=1.0,
                    value=0.7,
                    step=0.1,
                    info="বেশি = সৃজনশীল, কম = সুসংগত"
                )
                
                top_p_slider = gr.Slider(
                    label="Top P",
                    minimum=0.0,
                    maximum=1.0,
                    value=0.9,
                    step=0.05,
                    info="শব্দ নির্বাচন নিয়ন্ত্রণ"
                )
            
            submit_btn = gr.Button("✨ Generate Response", variant="primary", size="lg")
        
        with gr.Column():
            output_text = gr.Textbox(
                label="Response",
                lines=8,
                interactive=False
            )
    
    # Examples
    gr.Examples(
        examples=[
            ["বাংলা ভাষা কি?"],
            ["আমাকে একটা গল্প বলো"],
            ["পাইথন প্রোগ্রামিং কি?"],
        ],
        inputs=user_input,
    )
    
    # Button click handler
    submit_btn.click(
        fn=generate_response,
        inputs=[user_input, temp_slider, top_p_slider],
        outputs=output_text,
    )
    
    # Enter key handler
    user_input.submit(
        fn=generate_response,
        inputs=[user_input, temp_slider, top_p_slider],
        outputs=output_text,
    )

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