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

# Load the DeepSeek-R1-Distill-Qwen-1.5B-uncensored model
model_id = "thirdeyeai/DeepSeek-R1-Distill-Qwen-1.5B-uncensored"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.float16,  # Use float16 for efficiency
    low_cpu_mem_usage=True,
    device_map="auto"  # Automatically use available devices
)

def generate_text(prompt, max_length=100, temperature=0.7, top_p=0.9):
    """Generate text based on prompt"""
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    
    # Generate
    with torch.no_grad():
        generation_output = model.generate(
            input_ids=inputs.input_ids,
            attention_mask=inputs.attention_mask,
            max_length=len(inputs.input_ids[0]) + max_length,
            temperature=temperature,
            top_p=top_p,
            do_sample=True,
        )
    
    # Decode and return only the generated part
    generated_text = tokenizer.decode(generation_output[0], skip_special_tokens=True)
    return generated_text

# Create Gradio interface
demo = gr.Interface(
    fn=generate_text,
    inputs=[
        gr.Textbox(lines=5, placeholder="Enter your prompt here...", label="Prompt"),
        gr.Slider(minimum=10, maximum=500, value=100, step=10, label="Max Length"),
        gr.Slider(minimum=0.1, maximum=2.0, value=0.7, step=0.1, label="Temperature"),
        gr.Slider(minimum=0.1, maximum=1.0, value=0.9, step=0.05, label="Top-p")
    ],
    outputs=gr.Textbox(label="Generated Text"),
    title="DeepSeek-R1-Distill-Qwen-1.5B Demo",
    description="Enter a prompt to generate text with the DeepSeek-R1-Distill-Qwen-1.5B-uncensored model."
)

# Launch the app
demo.launch()