Dramatron: Hierarchical Story Generation System

The StoryGenerator follows a hierarchical generation process. You can use step() to generate new content at a specific level or complete() to extend existing content. You can also use rewrite() to modify existing content.

Generation Levels:

• Level 0: Title
• Level 1: Characters
• Level 2: Plot Synopsis (Scenes)
• Level 3: Place Descriptions
• Level 4: Dialogues

Key Methods:

• step(level, seed=None, idx=None): Generates new content for the specified level. For dialogues, provide idx to specify the scene.
• complete(level, seed=None, entity=None, sample_length=None): Continues/extends the current content at the specified level.
• rewrite(text, level, entity=None): Rewrites the content at the specified level. For places or dialogues, provide the entity (name or index) to target the specific item.
• get_story(): Returns the full story object.

Dramatron uses specific prompt prefixes to guide the LLM through different stages of storytelling (characters, scenes, settings, titles, and dialogue). You can create a custom_prefixes dictionary to override the default templates.

When writing custom prefixes, you must use the project's specific element markers and termination markers to ensure the LLM understands the structure:

Element Markers:

• LOGLINE_MARKER: Indicates the story logline.
• CHARACTER_MARKER: Marks a character definition.
• DESCRIPTION_MARKER: Marks a character or setting description.
• SCENES_MARKER: Marks the beginning of a scene list.
• PLACE_ELEMENT: Marks a location.
• PLOT_ELEMENT: Marks a plot point.
• BEAT_ELEMENT: Marks a specific story beat.
• DIALOG_MARKER: Marks the start of dialogue.
• TITLE_ELEMENT: Marks a title.
• SUMMARY_ELEMENT: Marks a story summary.

Termination Markers:

• STOP_MARKER: Used to signal the end of a specific element (like a character description).
• END_MARKER: Used to signal the end of a complete example block.

To use Google Gemini as the language model for Dramatron, initialize the GoogleAPI class. You will need a GOOGLE_API_KEY from Google AI Studio. The configuration includes the model name (e.g., gemini-1.5-pro), sampling parameters, and a system prompt that instructs the model to act as a writing assistant and follow specific storytelling formats.

Key configuration parameters for the config dictionary:

• language_api_name: Set to 'Gemini'.
• model_param: The selected GEMINI_MODEL_NAME.
• model_name: The selected GEMINI_MODEL_NAME.
• sampling: A dictionary containing prob, temp, and top_k values.
• max_retries: Number of retries for remote API calls.

import google.generativeai as genai

# Configuration setup
config = {}
config['language_api_name'] = 'Gemini'
config['model_param'] = GEMINI_MODEL_NAME
config['model_name'] = GEMINI_MODEL_NAME
config['max_retries'] = MAX_RETRIES
config['sample_length'] = SAMPLE_LENGTH
config['sampling'] = {
    'prob': SAMPLING_PROB,
    'temp': SAMPLING_TEMP,
    'top_k': GEMINI_SAMPLING_TOP_K
}

client = GoogleAPI(
    model_param=config['model_param'],
    model=config['model_name'],
    seed=DEFAULT_SEED,
    sample_length=config['sample_length'],
    max_retries=config['max_retries'],
    config_sampling=config['sampling']
)

To use OpenAI's Chat Completion API, initialize the OpenAIAPI class. This requires an OPENAI_API_KEY. The model_param used in the configuration corresponds to the CHATGPT_SYSTEM_PROMPT.

Key configuration parameters for the config dictionary:

• language_api_name: Set to 'OpenAI'.
• model_param: The CHATGPT_SYSTEM_PROMPT.
• model_name: The selected CHATGPT_MODEL_NAME (e.g., gpt-4-1106-preview).
• sampling: A dictionary containing prob, temp, frequency_penalty, and presence_penalty.

Note: Ensure os.environ['OPENAI_API_KEY'] is set before initializing.

import os
from openai import OpenAI

# Configuration setup
config = {}
config['language_api_name'] = 'OpenAI'
config['model_param'] = CHATGPT_SYSTEM_PROMPT
config['model_name'] = CHATGPT_MODEL_NAME
config['sampling'] = {
    'prob': SAMPLING_PROB,
    'temp': SAMPLING_TEMP,
    'frequency_penalty': CHATGPT_FREQUENCY_PENALTY,
    'presence_penalty': CHATGPT_PRESENCE_PENALTY
}

client = OpenAIAPI(
    model_param=config['model_param'],
    model=config['model_name'],
    seed=DEFAULT_SEED,
    sample_length=config['sample_length'],
    max_retries=config['max_retries'],
    config_sampling=config['sampling']
)

To use Mixtral models via the Groq API, initialize the GroqAPI class. This requires a GROQ_API_KEY set in the environment variables. The model_param serves as the system prompt for the chat completion.

Key configuration parameters for the config dictionary:

• language_api_name: Set to 'Groq' (implied by class usage).
• model_param: The GROQ_SYSTEM_PROMPT.
• model_name: The selected GROQ_MODEL_NAME (e.g., mixtral-8x7b-32768).
• sampling: A dictionary containing prob and temp values.

from groq import Groq
import os

# Configuration setup
config = {}
config['model_param'] = GROQ_SYSTEM_PROMPT
config['model_name'] = GROQ_MODEL_NAME
config['sampling'] = {
    'prob': SAMPLING_PROB,
    'temp': SAMPLING_TEMP
}

client = GroqAPI(
    model_param=config['model_param'],
    model=config['model_name'],
    seed=DEFAULT_SEED,
    sample_length=config['sample_length'],
    max_retries=config['max_retries'],
    config_sampling=config['sampling']
)