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

# Load the tokenizer and model
tokenizer = AutoTokenizer.from_pretrained("deepseek-ai/DeepSeek-V3")
model = AutoModelForCausalLM.from_pretrained(
    "deepseek-ai/DeepSeek-V3",
    torch_dtype=torch.float16,  # Use FP16 for efficiency
    device_map="auto"           # Automatically use GPU if available
)

# Function to generate text
def generate_text(prompt, max_length=100, temperature=0.7, top_k=50):
    # Tokenize the input prompt
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    
    # Generate text
    outputs = model.generate(
        inputs["input_ids"],
        max_length=max_length,
        temperature=temperature,
        top_k=top_k,
        do_sample=True
    )
    
    # Decode the generated text
    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return generated_text

# Gradio interface
def gradio_interface(prompt, max_length, temperature, top_k):
    try:
        output = generate_text(prompt, max_length, temperature, top_k)
        return output
    except Exception as e:
        return f"Error: {str(e)}"

# Custom CSS for a fancy theme
custom_css = """
.gradio-container {
    background: linear-gradient(135deg, #1e3c72, #2a5298);
    color: white;
    font-family: 'Arial', sans-serif;
}
h1 {
    color: #ffd700;
    text-align: center;
}
.description {
    color: #ffffff;
    text-align: center;
    font-size: 16px;
}
.input-label, .output-label {
    color: #ffffff;
}
.slider-label {
    color: #ffffff;
}
"""

# Create the Gradio app
iface = gr.Interface(
    fn=gradio_interface,
    inputs=[
        gr.Textbox(label="Enter your prompt", lines=3, placeholder="Once upon a time..."),
        gr.Slider(minimum=10, maximum=200, value=100, label="Max Length"),
        gr.Slider(minimum=0.1, maximum=1.0, value=0.7, label="Temperature"),
        gr.Slider(minimum=1, maximum=100, value=50, label="Top-k Sampling")
    ],
    outputs=gr.Textbox(label="Generated Text", lines=10),
    title="DreamWeaver AI",
    description="Crafting Stories, One Prompt at a Time. Generate text using the DeepSeek-V3 model. Adjust the parameters to control the output.",
    css=custom_css
)

# Launch the app
iface.launch()