import torch
import transformers
import gradio as gr

from transformers import AutoModelForCausalLM, GPT2Tokenizer

model_name = "DarwinAnim8or/GPT-NoSleep-v2"  
tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
model = AutoModelForCausalLM.from_pretrained(model_name)

# Handle padding
if tokenizer.pad_token_id is None:
    tokenizer.pad_token_id = tokenizer.eos_token_id

def generate_story(prompt, max_length=200, temp=0.3):
    """Generates a story continuation from a given prompt."""
    input_ids = tokenizer.encode(prompt, return_tensors="pt")

    # Set generation parameters (adjust for creativity)
    output = model.generate(
        input_ids,
        max_length=max_length,
        num_return_sequences=1,  # Generate a single story
        no_repeat_ngram_size=2,  # Prevent repetitive phrases
        do_sample=True,
        top_k=50,
        top_p=0.95,
        temperature=temp,      # Control randomness (higher = more creative)
    )

    # Decode the generated story
    story = tokenizer.decode(output[0], skip_special_tokens=True)
    return story

# Gradio Interface
with gr.Blocks() as demo:
    gr.Markdown("## 'NoSleep' Storyteller: Generate a story from a prompt!")
    prompt_input = gr.Textbox(label="Enter your story prompt:")
    story_output = gr.Textbox(label="Generated story:")
    max_length_slider = gr.Slider(minimum=50, maximum=500, value=200, step=10, label="Max Story Length")
    temp_slider = gr.Slider(minimum=0.1, maximum=1.0, value=0.3, step=0.1, label="Temperature (randomness)")
    generate_button = gr.Button("Generate Story")

    # Event handling
    generate_button.click(
        fn=generate_story, 
        inputs=[prompt_input, max_length_slider, temp_slider], 
        outputs=story_output
    )

# Launch the demo (customize the sharing options if desired)
demo.launch()