gpt_author.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# gpt-author\n",
        "By Matt Shumer (https://twitter.com/mattshumer_)\n",
        "\n",
        "Github repo: https://github.com/mshumer/gpt-author\n",
        "\n",
        "Generate an entire novel in minutes, and automatically package it as an e-book.\n",
        "\n",
        "To generate a book:\n",
        "1.  In the first cell, add in your OpenAI and Stability keys (see the first cell for instructions to get them).\n",
        "2.  Fill in the prompt, number of chapters, and writing style in the last cell.\n",
        "3. Run all the cells! After some time, your EPUB file should appear in the filesystem.\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "Ih8GSkb28daX"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install openai\n",
        "!pip install EbookLib\n",
        "\n",
        "import openai\n",
        "import os\n",
        "from ebooklib import epub\n",
        "import base64\n",
        "import os\n",
        "import requests\n",
        "\n",
        "openai.api_key = \"ENTER OPENAI KEY HERE\" # get it at https://platform.openai.com/\n",
        "stability_api_key = \"ENTER STABILITY KEY HERE\" # get it at https://beta.dreamstudio.ai/"
      ],
      "metadata": {
        "id": "HTGF4Bw6Hnuz"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def generate_cover_prompt(plot):\n",
        "    response = openai.ChatCompletion.create(\n",
        "        model=\"gpt-3.5-turbo-16k\",\n",
        "        messages=[\n",
        "            {\"role\": \"system\", \"content\": \"You are a creative assistant that writes a spec for the cover art of a book, based on the book's plot.\"},\n",
        "            {\"role\": \"user\", \"content\": f\"Plot: {plot}\\n\\n--\\n\\nDescribe the cover we should create, based on the plot. This should be two sentences long, maximum.\"}\n",
        "        ]\n",
        "    )\n",
        "    return response['choices'][0]['message']['content']\n",
        "\n",
        "\n",
        "def create_cover_image(plot):\n",
        "\n",
        "  plot = str(generate_cover_prompt(plot))\n",
        "\n",
        "  engine_id = \"stable-diffusion-xl-beta-v2-2-2\"\n",
        "  api_host = os.getenv('API_HOST', 'https://api.stability.ai')\n",
        "  api_key = stability_api_key\n",
        "\n",
        "  if api_key is None:\n",
        "      raise Exception(\"Missing Stability API key.\")\n",
        "\n",
        "  response = requests.post(\n",
        "      f\"{api_host}/v1/generation/{engine_id}/text-to-image\",\n",
        "      headers={\n",
        "          \"Content-Type\": \"application/json\",\n",
        "          \"Accept\": \"application/json\",\n",
        "          \"Authorization\": f\"Bearer {api_key}\"\n",
        "      },\n",
        "      json={\n",
        "          \"text_prompts\": [\n",
        "              {\n",
        "                  \"text\": plot\n",
        "              }\n",
        "          ],\n",
        "          \"cfg_scale\": 7,\n",
        "          \"clip_guidance_preset\": \"FAST_BLUE\",\n",
        "          \"height\": 768,\n",
        "          \"width\": 512,\n",
        "          \"samples\": 1,\n",
        "          \"steps\": 30,\n",
        "      },\n",
        "  )\n",
        "\n",
        "  if response.status_code != 200:\n",
        "      raise Exception(\"Non-200 response: \" + str(response.text))\n",
        "\n",
        "  data = response.json()\n",
        "\n",
        "  for i, image in enumerate(data[\"artifacts\"]):\n",
        "      with open(f\"/content/cover.png\", \"wb\") as f: # replace this if running locally, to where you store the cover file\n",
        "          f.write(base64.b64decode(image[\"base64\"]))"
      ],
      "metadata": {
        "id": "__nvOnmhPqwy"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def create_epub(title, author, chapters, cover_image_path='cover.png'):\n",
        "    book = epub.EpubBook()\n",
        "\n",
        "    # Set metadata\n",
        "    book.set_identifier('id123456')\n",
        "    book.set_title(title)\n",
        "    book.set_language('en')\n",
        "    book.add_author(author)\n",
        "\n",
        "    # Add cover image\n",
        "    with open(cover_image_path, 'rb') as cover_file:\n",
        "        cover_image = cover_file.read()\n",
        "    book.set_cover('cover.png', cover_image)\n",
        "\n",
        "    # Create chapters and add them to the book\n",
        "    epub_chapters = []\n",
        "    for i, chapter_dict in enumerate(chapters):\n",
        "        full_chapter_title = list(chapter_dict.keys())[0]\n",
        "        chapter_content = list(chapter_dict.values())[0]\n",
        "        if ' - ' in full_chapter_title:\n",
        "            chapter_title = full_chapter_title.split(' - ')[1]\n",
        "        else:\n",
        "            chapter_title = full_chapter_title\n",
        "\n",
        "        chapter_file_name = f'chapter_{i+1}.xhtml'\n",
        "        epub_chapter = epub.EpubHtml(title=chapter_title, file_name=chapter_file_name, lang='en')\n",
        "\n",
        "        # Add paragraph breaks\n",
        "        formatted_content = ''.join(f'<p>{paragraph.strip()}</p>' for paragraph in chapter_content.split('\\n') if paragraph.strip())\n",
        "\n",
        "        epub_chapter.content = f'<h1>{chapter_title}</h1>{formatted_content}'\n",
        "        book.add_item(epub_chapter)\n",
        "        epub_chapters.append(epub_chapter)\n",
        "\n",
        "\n",
        "    # Define Table of Contents\n",
        "    book.toc = (epub_chapters)\n",
        "\n",
        "    # Add default NCX and Nav files\n",
        "    book.add_item(epub.EpubNcx())\n",
        "    book.add_item(epub.EpubNav())\n",
        "\n",
        "    # Define CSS style\n",
        "    style = '''\n",
        "    @namespace epub \"http://www.idpf.org/2007/ops\";\n",
        "    body {\n",
        "        font-family: Cambria, Liberation Serif, serif;\n",
        "    }\n",
        "    h1 {\n",
        "        text-align: left;\n",
        "        text-transform: uppercase;\n",
        "        font-weight: 200;\n",
        "    }\n",
        "    '''\n",
        "\n",
        "    # Add CSS file\n",
        "    nav_css = epub.EpubItem(uid=\"style_nav\", file_name=\"style/nav.css\", media_type=\"text/css\", content=style)\n",
        "    book.add_item(nav_css)\n",
        "\n",
        "    # Create spine\n",
        "    book.spine = ['nav'] + epub_chapters\n",
        "\n",
        "    # Save the EPUB file\n",
        "    epub.write_epub(f'{title}.epub', book)"
      ],
      "metadata": {
        "id": "4w-tRWAVDf-V"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import openai\n",
        "import random\n",
        "import json\n",
        "import ast\n",
        "\n",
        "def print_step_costs(response, model):\n",
        "  input = response['usage']['prompt_tokens']\n",
        "  output = response['usage']['completion_tokens']\n",
        "\n",
        "  if model == \"gpt-4\" or model == \"gpt-4-0613\":\n",
        "    input_per_token = 0.00003\n",
        "    output_per_token = 0.00006\n",
        "  if model == \"gpt-3.5-turbo-16k\":\n",
        "    input_per_token = 0.000003\n",
        "    output_per_token = 0.000004\n",
        "  if model == \"gpt-4-32k-0613\" or model == \"gpt-4-32k\":\n",
        "    input_per_token = 0.00006\n",
        "    output_per_token = 0.00012\n",
        "  if model == \"gpt-3.5-turbo\" or model == \"gpt-3.5-turbo-0613\":\n",
        "    input_per_token = 0.0000015\n",
        "    output_per_token = 0.000002\n",
        "\n",
        "  input_cost = int(input) * input_per_token\n",
        "  output_cost = int(output) * output_per_token\n",
        "\n",
        "  total_cost = input_cost + output_cost\n",
        "  print('step cost:', total_cost)\n",
        "\n",
        "def generate_plots(prompt):\n",
        "    response = openai.ChatCompletion.create(\n",
        "        model=\"gpt-4-0613\",\n",
        "        messages=[\n",
        "            {\"role\": \"system\", \"content\": \"You are a creative assistant that generates engaging fantasy novel plots.\"},\n",
        "            {\"role\": \"user\", \"content\": f\"Generate 10 fantasy novel plots based on this prompt: {prompt}\"}\n",
        "        ]\n",
        "    )\n",
        "\n",
        "    print_step_costs(response, \"gpt-4-0613\")\n",
        "\n",
        "    return response['choices'][0]['message']['content'].split('\\n')\n",
        "\n",
        "def select_most_engaging(plots):\n",
        "    response = openai.ChatCompletion.create(\n",
        "        model=\"gpt-4-0613\",\n",
        "        messages=[\n",
        "            {\"role\": \"system\", \"content\": \"You are an expert in writing fantastic fantasy novel plots.\"},\n",
        "            {\"role\": \"user\", \"content\": f\"Here are a number of possible plots for a new novel: {plots}\\n\\n--\\n\\nNow, write the final plot that we will go with. It can be one of these, a mix of the best elements of multiple, or something completely new and better. The most important thing is the plot should be fantastic, unique, and engaging.\"}\n",
        "        ]\n",
        "    )\n",
        "\n",
        "    print_step_costs(response, \"gpt-4-0613\")\n",
        "\n",
        "    return response['choices'][0]['message']['content']\n",
        "\n",
        "def improve_plot(plot):\n",
        "    response = openai.ChatCompletion.create(\n",
        "        model=\"gpt-4-0613\",\n",
        "        messages=[\n",
        "            {\"role\": \"system\", \"content\": \"You are an expert in improving and refining story plots.\"},\n",
        "            {\"role\": \"user\", \"content\": f\"Improve this plot: {plot}\"}\n",
        "        ]\n",
        "    )\n",
        "\n",
        "    print_step_costs(response, \"gpt-4-0613\")\n",
        "\n",
        "    return response['choices'][0]['message']['content']\n",
        "\n",
        "def get_title(plot):\n",
        "    response = openai.ChatCompletion.create(\n",
        "        model=\"gpt-3.5-turbo-16k\",\n",
        "        messages=[\n",
        "            {\"role\": \"system\", \"content\": \"You are an expert writer.\"},\n",
        "            {\"role\": \"user\", \"content\": f\"Here is the plot: {plot}\\n\\nWhat is the title of this book? Just respond with the title, do nothing else.\"}\n",
        "        ]\n",
        "    )\n",
        "\n",
        "    print_step_costs(response, \"gpt-3.5-turbo-16k\")\n",
        "\n",
        "    return response['choices'][0]['message']['content']\n",
        "\n",
        "def write_first_chapter(plot, first_chapter_title, writing_style):\n",
        "    response = openai.ChatCompletion.create(\n",
        "        model=\"gpt-4-0613\",\n",
        "        messages=[\n",
        "            {\"role\": \"system\", \"content\": \"You are a world-class fantasy writer.\"},\n",
        "            {\"role\": \"user\", \"content\": f\"Here is the high-level plot to follow: {plot}\\n\\nWrite the first chapter of this novel: `{first_chapter_title}`.\\n\\nMake it incredibly unique, engaging, and well-written.\\n\\nHere is a description of the writing style you should use: `{writing_style}`\\n\\nInclude only the chapter text. There is no need to rewrite the chapter name.\"}\n",
        "        ]\n",
        "    )\n",
        "\n",
        "    print_step_costs(response, \"gpt-4-0613\")\n",
        "\n",
        "    improved_response = openai.ChatCompletion.create(\n",
        "        model=\"gpt-4-32k-0613\",\n",
        "        messages=[\n",
        "            {\"role\": \"system\", \"content\": \"You are a world-class fantasy writer. Your job is to take your student's rough initial draft of the first chapter of their fantasy novel, and rewrite it to be significantly better, with much more detail.\"},\n",
        "            {\"role\": \"user\", \"content\": f\"Here is the high-level plot you asked your student to follow: {plot}\\n\\nHere is the first chapter they wrote: {response['choices'][0]['message']['content']}\\n\\nNow, rewrite the first chapter of this novel, in a way that is far superior to your student's chapter. It should still follow the exact same plot, but it should be far more detailed, much longer, and more engaging. Here is a description of the writing style you should use: `{writing_style}`\"}\n",
        "        ]\n",
        "    )\n",
        "\n",
        "    print_step_costs(response, \"gpt-4-32k-0613\")\n",
        "\n",
        "    return improved_response['choices'][0]['message']['content']\n",
        "\n",
        "def write_chapter(previous_chapters, plot, chapter_title):\n",
        "    try:\n",
        "      i = random.randint(1,2242)\n",
        "      # write_to_file(f'write_chapter_{i}', f\"Plot: {plot}, Previous Chapters: {previous_chapters}\\n\\n--\\n\\nWrite the next chapter of this novel, following the plot and taking in the previous chapters as context. Here is the plan for this chapter: {chapter_title}\\n\\nWrite it beautifully. Include only the chapter text. There is no need to rewrite the chapter name.\")\n",
        "      response = openai.ChatCompletion.create(\n",
        "        model=\"gpt-4-0613\",\n",
        "        messages=[\n",
        "            {\"role\": \"system\", \"content\": \"You are a world-class fantasy writer.\"},\n",
        "            {\"role\": \"user\", \"content\": f\"Plot: {plot}, Previous Chapters: {previous_chapters}\\n\\n--\\n\\nWrite the next chapter of this novel, following the plot and taking in the previous chapters as context. Here is the plan for this chapter: {chapter_title}\\n\\nWrite it beautifully. Include only the chapter text. There is no need to rewrite the chapter name.\"}\n",
        "        ]\n",
        "    )\n",
        "\n",
        "      print_step_costs(response, \"gpt-4-0613\")\n",
        "\n",
        "      return response['choices'][0]['message']['content']\n",
        "    except:\n",
        "      response = openai.ChatCompletion.create(\n",
        "        model=\"gpt-4-32k-0613\",\n",
        "        messages=[\n",
        "            {\"role\": \"system\", \"content\": \"You are a world-class fantasy writer.\"},\n",
        "            {\"role\": \"user\", \"content\": f\"Plot: {plot}, Previous Chapters: {previous_chapters}\\n\\n--\\n\\nWrite the next chapter of this novel, following the plot and taking in the previous chapters as context. Here is the plan for this chapter: {chapter_title}\\n\\nWrite it beautifully. Include only the chapter text. There is no need to rewrite the chapter name.\"}\n",
        "        ]\n",
        "    )\n",
        "\n",
        "      print_step_costs(response, \"gpt-4-32k-0613\")\n",
        "\n",
        "      return response['choices'][0]['message']['content']\n",
        "\n",
        "\n",
        "def generate_storyline(prompt, num_chapters):\n",
        "    print(\"Generating storyline with chapters and high-level details...\")\n",
        "    json_format = \"\"\"[{\"Chapter CHAPTER_NUMBER_HERE - CHAPTER_TITLE_GOES_HERE\": \"CHAPTER_OVERVIEW_AND_DETAILS_GOES_HERE\"}, ...]\"\"\"\n",
        "    response = openai.ChatCompletion.create(\n",
        "        model=\"gpt-4-0613\",\n",
        "        messages=[\n",
        "            {\"role\": \"system\", \"content\": \"You are a world-class fantasy writer. Your job is to write a detailed storyline, complete with chapters, for a fantasy novel. Don't be flowery -- you want to get the message across in as few words as possible. But those words should contain lots of information.\"},\n",
        "            {\"role\": \"user\", \"content\": f'Write a fantastic storyline with {num_chapters} chapters and high-level details based on this plot: {prompt}.\\n\\nDo it in this list of dictionaries format {json_format}'}\n",
        "        ]\n",
        "    )\n",
        "\n",
        "    print_step_costs(response, \"gpt-4-0613\")\n",
        "\n",
        "    improved_response = openai.ChatCompletion.create(\n",
        "        model=\"gpt-4-0613\",\n",
        "        messages=[\n",
        "            {\"role\": \"system\", \"content\": \"You are a world-class fantasy writer. Your job is to take your student's rough initial draft of the storyline of a fantasy novel, and rewrite it to be significantly better.\"},\n",
        "            {\"role\": \"user\", \"content\": f\"Here is the draft storyline they wrote: {response['choices'][0]['message']['content']}\\n\\nNow, rewrite the storyline, in a way that is far superior to your student's version. It should have the same number of chapters, but it should be much improved in as many ways as possible. Remember to do it in this list of dictionaries format {json_format}\"}\n",
        "        ]\n",
        "    )\n",
        "\n",
        "    print_step_costs(improved_response, \"gpt-4-0613\")\n",
        "\n",
        "    return improved_response['choices'][0]['message']['content']\n",
        "\n",
        "\n",
        "def write_to_file(prompt, content):\n",
        "\n",
        "    # Create a directory for the prompts if it doesn't exist\n",
        "    if not os.path.exists('prompts'):\n",
        "        os.mkdir('prompts')\n",
        "\n",
        "    # Replace invalid characters for filenames\n",
        "    valid_filename = ''.join(c for c in prompt if c.isalnum() or c in (' ', '.', '_')).rstrip()\n",
        "    file_path = f'prompts/{valid_filename}.txt'\n",
        "\n",
        "    with open(file_path, 'w', encoding='utf-8') as f:\n",
        "        f.write(content)\n",
        "\n",
        "    print(f'Output for prompt \"{prompt}\" has been written to {file_path}\\n')\n",
        "\n",
        "\n",
        "def write_fantasy_novel(prompt, num_chapters, writing_style):\n",
        "    plots = generate_plots(prompt)\n",
        "\n",
        "    best_plot = select_most_engaging(plots)\n",
        "\n",
        "    improved_plot = improve_plot(best_plot)\n",
        "\n",
        "\n",
        "    title = get_title(improved_plot)\n",
        "\n",
        "    storyline = generate_storyline(improved_plot, num_chapters)\n",
        "    chapter_titles = ast.literal_eval(storyline)\n",
        "\n",
        "\n",
        "    novel = f\"Storyline:\\n{storyline}\\n\\n\"\n",
        "\n",
        "    first_chapter = write_first_chapter(storyline, chapter_titles[0], writing_style.strip())\n",
        "    novel += f\"Chapter 1:\\n{first_chapter}\\n\"\n",
        "    chapters = [first_chapter]\n",
        "\n",
        "    for i in range(num_chapters - 1):\n",
        "        print(f\"Writing chapter {i+2}...\") # + 2 because the first chapter was already added\n",
        "        chapter = write_chapter(novel, storyline, chapter_titles[i+1])\n",
        "        novel += f\"Chapter {i+2}:\\n{chapter}\\n\"\n",
        "        chapters.append(chapter)\n",
        "\n",
        "    return novel, title, chapters, chapter_titles"
      ],
      "metadata": {
        "id": "KEClRDpDK7VW"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Example usage:\n",
        "prompt = \"A kingdom hidden deep in the forest, where every tree is a portal to another world.\"\n",
        "num_chapters = 20\n",
        "writing_style = \"Clear and easily understandable, similar to a young adult novel. Highly descriptive and sometimes long-winded.\"\n",
        "novel, title, chapters, chapter_titles = write_fantasy_novel(prompt, num_chapters, writing_style)\n",
        "\n",
        "# Replace chapter descriptions with body text in chapter_titles\n",
        "for i, chapter in enumerate(chapters):\n",
        "    chapter_number_and_title = list(chapter_titles[i].keys())[0]\n",
        "    chapter_titles[i] = {chapter_number_and_title: chapter}\n",
        "\n",
        "# Create the cover\n",
        "create_cover_image(str(chapter_titles))\n",
        "\n",
        "# Create the EPUB file\n",
        "create_epub(title, 'AI', chapter_titles, '/content/cover.png')"
      ],
      "metadata": {
        "id": "IZeHy7qyACCb"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}