import gradio as gr
from transformers import pipeline
import torch

# Initialize the text generation pipeline with the model
generator = pipeline(
    "text-generation", 
    model="thirdeyeai/DeepSeek-R1-Distill-Qwen-1.5B-uncensored",
    torch_dtype=torch.float16,
    device_map="auto"
)

def generate_text(prompt, max_length=100, temperature=0.7, top_p=0.9):
    """Generate text based on prompt using the pipeline"""
    # Calculate max_new_tokens from max_length
    # This is approximate as token count doesn't directly map to character count
    max_new_tokens = max_length // 4  # rough estimate of 4 chars per token
    
    # Generate text
    response = generator(
        prompt,
        max_new_tokens=max_new_tokens,
        temperature=temperature,
        top_p=top_p,
        do_sample=True,
        return_full_text=True
    )
    
    # Extract the generated text from the response
    generated_text = response[0]['generated_text']
    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 (approx. characters)"),
        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()