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

# Model ID
MODEL_ID = "mx-llms/Lychee-GPT-9B"

# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)

# Load model
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.float32,
    device_map="cpu",
    trust_remote_code=True,
)

def generate_text(prompt, max_length=256, temperature=0.7, top_p=0.9):
    """Generate text using Lychee-GPT-9B"""
    try:
        # Tokenize input
        inputs = tokenizer(prompt, return_tensors="pt")
        
        # Generate
        with torch.no_grad():
            output = model.generate(
                inputs["input_ids"],
                max_new_tokens=max_length,
                temperature=temperature,
                top_p=top_p,
                do_sample=True,
                pad_token_id=tokenizer.eos_token_id,
            )
        
        # Decode output
        response = tokenizer.decode(output[0], skip_special_tokens=True)
        
        # Remove input prompt from response
        if prompt in response:
            response = response.replace(prompt, "", 1).strip()
        
        return response
    except Exception as e:
        return f"Error: {str(e)}"

# Create Gradio interface
with gr.Blocks(title="Lychee-GPT-9B") as demo:
    gr.Markdown("""
    # 🎉 Lychee-GPT-9B Demo
    
    আপনার নিজস্ব LLM Model!
    
    ⏱️ **নোট:** 
    - প্রথম load: 2-3 মিনিট
    - Response time: 30-90 সেকেন্ড (CPU তে)
    """)
    
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(
                label="প্রশ্ন/Prompt",
                placeholder="কিছু লিখুন...",
                lines=3
            )
            
            with gr.Row():
                max_len = gr.Slider(
                    label="Max Length",
                    minimum=10,
                    maximum=512,
                    value=256,
                    step=10
                )
            
            with gr.Row():
                temp = gr.Slider(
                    label="Temperature",
                    minimum=0.0,
                    maximum=1.0,
                    value=0.7,
                    step=0.1
                )
                
                top_p = gr.Slider(
                    label="Top P",
                    minimum=0.0,
                    maximum=1.0,
                    value=0.9,
                    step=0.05
                )
            
            submit_btn = gr.Button("Generate", variant="primary", size="lg")
        
        with gr.Column():
            output = gr.Textbox(
                label="Response",
                lines=8,
                interactive=False
            )
    
    # Examples
    gr.Examples(
        examples=[
            ["বাংলা ভাষা সম্পর্কে বলুন"],
            ["পাইথন প্রোগ্রামিং কি?"],
            ["একটি সংক্ষিপ্ত গল্প বলুন"],
        ],
        inputs=prompt,
    )
    
    # Connect button
    submit_btn.click(
        fn=generate_text,
        inputs=[prompt, max_len, temp, top_p],
        outputs=output
    )
    
    # Allow Enter key
    prompt.submit(
        fn=generate_text,
        inputs=[prompt, max_len, temp, top_p],
        outputs=output
    )

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