File size: 3,377 Bytes
6c48899 a86d089 79cf13e 6c48899 f844445 6c48899 f844445 6c48899 f844445 6c48899 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 |
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
# Model setup
#model_name = "EleutherAI/gpt-neo-125M" # Lightweight model for Spaces
model_name = "EleutherAI/gpt-neo-1.3B" # A bit better model
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name).to("cpu") # CPU-friendly for Spaces
# Text generation function
def generate_text(prompt, max_length=100, temperature=0.7, top_p=0.9):
inputs = tokenizer(prompt, return_tensors="pt")
outputs = model.generate(
**inputs,
max_length=max_length,
temperature=temperature,
top_p=top_p,
do_sample=True,
pad_token_id=tokenizer.eos_token_id
)
return tokenizer.decode(outputs[0], skip_special_tokens=True)
# Global variables to store hierarchical content
global_synopsis = ""
global_chapters = ""
# Generation functions
def generate_synopsis(topic):
global global_synopsis
prompt = f"Write a brief synopsis for a story about {topic}: "
global_synopsis = generate_text(prompt, max_length=100)
return global_synopsis
def generate_chapters():
global global_synopsis, global_chapters
if not global_synopsis:
return "Please generate a synopsis first."
prompt = f'''Based on this synopsis for a book: {global_synopsis}. Divide the story into 4 chapters with brief descriptions for each.
Enumerate every chapter created followed by its description and make the first chapter sound like an introduction and the last may sound as the epilogue'''
global_chapters = generate_text(prompt, max_length=2000)
return global_chapters
def expand_chapter(chapter_number):
global global_chapters
if not global_chapters:
return "Please generate chapters first."
chapters = global_chapters.split("\n")
if chapter_number <= 0 or chapter_number > len(chapters):
return f"Select a number between 1 and {len(chapters)}."
prompt = f'''Knowing this synopsis for a book: {global_synopsis}. Expand and describe Chapter {chapter_number}
in more detail, the title and current brief description of this chapter is: {chapters[chapter_number - 1]}'''
return generate_text(prompt, max_length=200)
# Gradio interface
with gr.Blocks() as demo:
gr.Markdown("## AI Hierarchical Story Generator")
with gr.Tab("Generate Synopsis"):
topic_input = gr.Textbox(label="Enter the story's main topic")
synopsis_output = gr.Textbox(label="Generated Synopsis", interactive=False)
synopsis_button = gr.Button("Generate Synopsis")
with gr.Tab("Generate Chapters"):
chapters_output = gr.Textbox(label="Generated Chapters", interactive=False)
chapters_button = gr.Button("Generate Chapters")
with gr.Tab("Expand Chapter"):
chapter_input = gr.Number(label="Chapter Number", precision=0)
chapter_detail_output = gr.Textbox(label="Expanded Chapter", interactive=False)
chapter_button = gr.Button("Expand Chapter")
# Connect functions to UI
synopsis_button.click(generate_synopsis, inputs=topic_input, outputs=synopsis_output)
chapters_button.click(generate_chapters, outputs=chapters_output)
chapter_button.click(expand_chapter, inputs=chapter_input, outputs=chapter_detail_output)
# Launch the app
demo.launch() |