#!/usr/bin/env python3
"""
Gradient.com এ Lychee-GPT-9B চালানোর জন্য script
Gradient Notebook এ এটা run করুন
"""

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

print("🚀 Installing dependencies...")
os.system("pip install -q torch transformers gradio accelerate safetensors")

print("📥 Loading Lychee-GPT-9B model...")
print("⏱️ এটা 5-10 মিনিট লাগতে পারে...")

MODEL_ID = "deepreinforce-ai/Ornith-1.0-9B"

try:
    # 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="auto",  # Gradient automatically optimizes
        trust_remote_code=True,
    )
    
    model.eval()
    print("✅ Model loaded successfully!")
    
except Exception as e:
    print(f"❌ Error loading model: {e}")
    model = None
    tokenizer = None

def generate_text(prompt, max_length=256, temperature=0.7, top_p=0.9):
    """Generate text using Lychee-GPT-9B"""
    if model is None or tokenizer is None:
        return "❌ Model failed to load"
    
    try:
        inputs = tokenizer(prompt, return_tensors="pt")
        
        with torch.no_grad():
            output = model.generate(
                inputs["input_ids"],
                max_new_tokens=int(max_length),
                temperature=float(temperature),
                top_p=float(top_p),
                do_sample=True,
                pad_token_id=tokenizer.eos_token_id,
            )
        
        response = tokenizer.decode(output[0], skip_special_tokens=True)
        
        if prompt in response:
            response = response.replace(prompt, "", 1).strip()
        
        return response if response else "No response generated"
        
    except Exception as e:
        return f"❌ Error: {str(e)}"

# Create Gradio interface
print("\n🎨 Creating Gradio interface...")

with gr.Blocks(title="Lychee-GPT-9B", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 🎉 Lychee-GPT-9B Demo
    
    আপনার নিজস্ব LLM Model - Gradient এ চলছে!
    
    ⏱️ **প্রথম response একটু slow হতে পারে (model warming up)**
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            prompt = gr.Textbox(
                label="প্রশ্ন/Prompt",
                placeholder="কিছু লিখুন...",
                lines=4,
                info="আপনার প্রশ্ন বা prompt দিন"
            )
            
            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(scale=1):
            output = gr.Textbox(
                label="Response",
                lines=10,
                interactive=False,
            )
    
    # Examples
    gr.Examples(
        examples=[
            ["বাংলা ভাষা সম্পর্কে বলুন"],
            ["পাইথন প্রোগ্রামিং কি?"],
            ["একটি সংক্ষিপ্ত গল্প বলুন"],
            ["কৃত্রিম বুদ্ধিমত্তা কি?"],
        ],
        inputs=prompt,
        label="উদাহরণ প্রশ্ন"
    )
    
    # 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,
    )

print("\n🌐 Launching Gradio interface...")
print("=" * 50)
print("Space URL will appear below 👇")
print("=" * 50)

# Launch
demo.launch(
    server_name="0.0.0.0",
    server_port=7860,
    share=True,
    show_error=True,
)