A quick hack of a python script to endlessly generate a story using koboldcpp as a backend AI text generator.

story_writer.py

import requests
import sys
import json
from pprint import pprint

# Requirements:
# * Koboldcpp (or compatible api)
#
# Setup:
# conda create -n creator python=3.10
# conda activate creator
# pip install requests
#
# Example CLI:
# python story_writer.py story.txt
#
# sample ideas for the story: Make it a classic fantasy fiction story in the style of Tolkien or Dungeons and Dragons, use the Hero Takes a Journey trope, main character should be highly rememberable and possess magical aptitudes.

# Consider changing `api_url` if needed and the `parameters` in `kobold_generate_text`

max_generated_length = 2048
context_len = 4096
api_url = 'http://localhost:5001/api/v1/generate'

# Setup for alpaca prompts
prompt_start = "### Instruction: "
prompt_end = "\n### Response: "
creator_personality = """The Writer, known by his pen name 'Inkwell', is a man of meticulous nature, with an unyielding dedication to the craft of fiction. He stands at a height of 6 feet, with dark, wavy hair that cascades over his forehead, often tousled due to the countless hours spent hunched over his desk. His eyes are deep and expressive, revealing the intensity of his thoughts and emotions, while his sharp features suggest a mind honed by years of introspection and observation.
Dressed simply in earthy tones, he exudes an air of understated elegance, allowing his work to speak volumes about his personality. His fingers, calloused from hours of writing, dance gracefully across the keys of his typewriter, as he weaves tales that captivate and enthrall his readers. His attention to detail borders on obsession, as he carefully crafts each sentence, ensuring that every word carries its intended weight and impact.
Inkwell is a man of strong discipline, rising early each day to capture the essence of dawn's inspiration before the world awakens. He approaches his writing with a clear and focused mind, allowing his ideas to flow unhindered onto the page. His vocabulary is vast, enabling him to paint vivid images and evoke powerful emotions through his words.
Despite his solitary profession, Inkwell is far from reclusive. He is well read, devouring books from various genres, fueling his own creativity and broadening his perspective. His strong creative urges compel him to explore new ideas and themes, constantly challenging himself to push the boundaries of his imagination.
Quirky and eccentric, he often paces around his study, muttering dialogues between characters or contemplating plot twists. A true artist at heart, Inkwell lives for the moments when his stories come alive on the page, transporting readers to other worlds and realities.
"""
creator_instruction = """As "The Writer", you are tasked with writing more parts of the following story. Specifically, The Writer should try to incorporate the following into their writing:
"""

text_to_token_estimate = 3.8 #number of chars per token on average, used to estimate budget


def kobold_generate_text(prompt):
    """Takes a text prompt and sends it to the KoboldCpp backend and returns a generated string or an empty string on failure."""
    parameters = {
        "max_context_length": context_len,
        "max_length": max_generated_length,
        "trim_stop": False,
        "prompt": prompt,
        "quiet": False,
        "rep_pen": 1.05,
        "rep_pen_range": 1024,
        "rep_pen_slope": 1,
        "temperature": 1.0,
        "min_p": 0.1,
    }
    response = requests.post(api_url, data=json.dumps(parameters))
    if response.status_code == 200:
        return response.json()['results'][0]['text'] # returns generated text by given prompt
    else:
        print("Error attempting to generate text. Response: " + str(response))
        return ""
    

def generate_story_block(feature_request, past_story):
    """Generates the story, continuing from where it let off."""

    prompt = prompt_start + creator_personality + creator_instruction + feature_request + "\n" 

    if len(past_story) < 1:
        prompt += "The story has not started yet, so start at the beginning of the tale.\n"
    else:
        # add the last part of the story, filling up the remaining budget for the prompt
        # leaving enough room for the max generated length
        prompt_len= len(prompt)
        total_prompt_budget_in_chars = (context_len - max_generated_length) * text_to_token_estimate
        remaining_budget_in_chars = total_prompt_budget_in_chars - prompt_len
        prompt += past_story[-int(remaining_budget_in_chars):]

    prompt += prompt_end 
    return kobold_generate_text(prompt)


def main(story_filename):
    print("\nWhat should The Writer try to incorporate into their writing? Give all the details here and an endless story will start.")
    feature_request = input(">> ")

    story = ""
    while True:
        try:
            print("\nGenerating...\n")
            new_story_chunk = generate_story_block(feature_request, story)
            print(new_story_chunk)
            print("\n")

            with open(story_filename, 'a') as story_file:
                story_file.write(new_story_chunk)

            story += new_story_chunk
        except Exception as e:
            print('An error occurred:', str(e))
            sys.exit(1)

if __name__=="__main__":
    if len(sys.argv) < 2:
        print("Usage: python " + sys.argv[0] + " <story_filename>")
        sys.exit(1)
    main(sys.argv[1])