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

# Load your model
print("🚀 Loading Starlight Mini...")

MODEL_NAME = "arthu1/starlight-mini"

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    torch_dtype=torch.float16,
    device_map="auto",
    low_cpu_mem_usage=True
)

print("✅ Model loaded!")

def generate_text(
    prompt: str,
    max_tokens: int = 100,
    temperature: float = 0.7,
    top_p: float = 0.9
) -> str:
    """Generate text from prompt"""
    if not prompt:
        return "Error: Please provide a prompt"
    
    try:
        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
        
        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=min(max_tokens, 500),
                temperature=temperature,
                top_p=top_p,
                do_sample=True,
                pad_token_id=tokenizer.eos_token_id
            )
        
        result = tokenizer.decode(outputs[0], skip_special_tokens=True)
        return result
    
    except Exception as e:
        return f"Error: {str(e)}"

# Create Gradio interface
with gr.Blocks(title="Starlight Mini API") as demo:
    gr.Markdown("# 🌟 Starlight Mini API\n**Free AI Text Generation - 8B Parameters**")
    
    with gr.Tab("💬 Chat"):
        with gr.Row():
            with gr.Column():
                prompt_input = gr.Textbox(label="Prompt", placeholder="Enter your prompt...", lines=5)
                
                with gr.Accordion("Settings", open=False):
                    max_tokens_slider = gr.Slider(10, 500, 100, step=10, label="Max Tokens")
                    temperature_slider = gr.Slider(0.1, 2.0, 0.7, step=0.1, label="Temperature")
                    top_p_slider = gr.Slider(0.1, 1.0, 0.9, step=0.05, label="Top P")
                
                generate_btn = gr.Button("✨ Generate", variant="primary")
            
            with gr.Column():
                output_text = gr.Textbox(label="Output", lines=10)
        
        gr.Examples(
            [["Once upon a time"], ["Explain AI:"], ["Code a function:"]], 
            inputs=[prompt_input]
        )
    
    with gr.Tab("📡 API Docs"):
        gr.Markdown("""
        ## API Usage
        
        **Python:**
        ```python
        from gradio_client import Client
        client = Client("arthu1/Exquisite-Starlight")
        result = client.predict("Hello!", 100, 0.7, 0.9, api_name="/predict")
        ```
        
        **cURL:**
        ```bash
        curl -X POST https://arthu1-exquisite-starlight.hf.space/api/predict \\
          -H "Content-Type: application/json" \\
          -d '{"data": ["Hello!", 100, 0.7, 0.9]}'
        ```
        """)
    
    generate_btn.click(
        fn=generate_text,
        inputs=[prompt_input, max_tokens_slider, temperature_slider, top_p_slider],
        outputs=output_text
    )

demo.launch(show_api=True, share=True)