Claude_Author.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyNxGwAz5u58p+OkkO/CDte3",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/mshumer/gpt-author/blob/main/Claude_Author.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-sATA0DDtule"
      },
      "outputs": [],
      "source": [
        "!pip install EbookLib\n",
        "\n",
        "import time\n",
        "import re\n",
        "import os\n",
        "from ebooklib import epub\n",
        "import base64\n",
        "import requests\n",
        "import json\n",
        "\n",
        "ANTHROPIC_API_KEY = \"YOUR KEY HERE\"\n",
        "stability_api_key = \"YOUR KEY HERE\" # get it at https://beta.dreamstudio.ai/\n",
        "\n",
        "def remove_first_line(test_string):\n",
        "    if test_string.startswith(\"Here\") and test_string.split(\"\\n\")[0].strip().endswith(\":\"):\n",
        "        return re.sub(r'^.*\\n', '', test_string, count=1)\n",
        "    return test_string\n",
        "\n",
        "def generate_text(prompt, model=\"claude-3-haiku-20240307\", max_tokens=2000, temperature=0.7):\n",
        "    headers = {\n",
        "        \"x-api-key\": ANTHROPIC_API_KEY,\n",
        "        \"anthropic-version\": \"2023-06-01\",\n",
        "        \"content-type\": \"application/json\"\n",
        "    }\n",
        "    data = {\n",
        "        \"model\": model,\n",
        "        \"max_tokens\": max_tokens,\n",
        "        \"temperature\": temperature,\n",
        "        \"system\": \"You are a world-class author. Write the requested content with great skill and attention to detail.\",\n",
        "        \"messages\": [{\"role\": \"user\", \"content\": prompt}],\n",
        "    }\n",
        "    response = requests.post(\"https://api.anthropic.com/v1/messages\", headers=headers, json=data)\n",
        "    response_text = response.json()['content'][0]['text']\n",
        "    return response_text.strip()\n",
        "\n",
        "def generate_cover_prompt(plot):\n",
        "    response = generate_text(f\"Plot: {plot}\\n\\n--\\n\\nDescribe the cover we should create, based on the plot. This should be two sentences long, maximum.\")\n",
        "    return response\n",
        "\n",
        "def generate_title(plot):\n",
        "    response = generate_text(f\"Here is the plot for the book: {plot}\\n\\n--\\n\\nRespond with a great title for this book. Only respond with the title, nothing else is allowed.\")\n",
        "    return remove_first_line(response)\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\"]))\n",
        "\n",
        "def generate_chapter_title(chapter_content):\n",
        "    response = generate_text(f\"Chapter Content:\\n\\n{chapter_content}\\n\\n--\\n\\nGenerate a concise and engaging title for this chapter based on its content. Respond with the title only, nothing else.\")\n",
        "    return remove_first_line(response)\n",
        "\n",
        "def create_epub(title, author, chapters, cover_image_path='cover.png'):\n",
        "    book = epub.EpubBook()\n",
        "    # Set metadata\n",
        "    book.set_identifier('id123456')\n",
        "    book.set_title(title)\n",
        "    book.set_language('en')\n",
        "    book.add_author(author)\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",
        "    # Create chapters and add them to the book\n",
        "    epub_chapters = []\n",
        "    for i, chapter_content in enumerate(chapters):\n",
        "        chapter_title = generate_chapter_title(chapter_content)\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",
        "        # Add paragraph breaks\n",
        "        formatted_content = ''.join(f'<p>{paragraph.strip()}</p>' for paragraph in chapter_content.split('\\n') if paragraph.strip())\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)\n",
        "\n",
        "\n",
        "def generate_book(writing_style, book_description, num_chapters):\n",
        "    print(\"Generating plot outline...\")\n",
        "    plot_prompt = f\"Create a detailed plot outline for a {num_chapters}-chapter book in the {writing_style} style, based on the following description:\\n\\n{book_description}\\n\\nEach chapter should be at least 10 pages long.\"\n",
        "    plot_outline = generate_text(plot_prompt)\n",
        "    print(\"Plot outline generated.\")\n",
        "\n",
        "    chapters = []\n",
        "    for i in range(num_chapters):\n",
        "        print(f\"Generating chapter {i+1}...\")\n",
        "        chapter_prompt = f\"Previous Chapters:\\n\\n{' '.join(chapters)}\\n\\nWriting style: `{writing_style}`\\n\\nPlot Outline:\\n\\n{plot_outline}\\n\\nWrite chapter {i+1} of the book, ensuring it follows the plot outline and builds upon the previous chapters. The chapter should be at least 256 paragraphs long... we're going for lengthy yet exciting chapters here.\"\n",
        "        chapter = generate_text(chapter_prompt, max_tokens=4000)\n",
        "        chapters.append(remove_first_line(chapter))\n",
        "        print(f\"Chapter {i+1} generated.\")\n",
        "        time.sleep(1)  # Add a short delay to avoid hitting rate limits\n",
        "\n",
        "    print(\"Compiling the book...\")\n",
        "    book = \"\\n\\n\".join(chapters)\n",
        "    print(\"Book generated!\")\n",
        "\n",
        "    return plot_outline, book, chapters\n",
        "\n",
        "# User input\n",
        "writing_style = input(\"Enter the desired writing style: \")\n",
        "book_description = input(\"Enter a high-level description of the book: \")\n",
        "num_chapters = int(input(\"Enter the number of chapters: \"))\n",
        "\n",
        "# Generate the book\n",
        "plot_outline, book, chapters = generate_book(writing_style, book_description, num_chapters)\n",
        "\n",
        "title = generate_title(plot_outline)\n",
        "\n",
        "# Save the book to a file\n",
        "with open(f\"{title}.txt\", \"w\") as file:\n",
        "    file.write(book)\n",
        "\n",
        "create_cover_image(plot_outline)\n",
        "\n",
        "# Create the EPUB file\n",
        "create_epub(title, 'AI', chapters, '/content/cover.png')\n",
        "\n",
        "print(f\"Book saved as '{title}.txt'.\")"
      ]
    }
  ]
}