author / app.py
Omnibus's picture
Update app.py
97259b9 verified
import gradio as gr
import time
import re
import os
from ebooklib import epub
import random
import base64
import requests
import json
from huggingface_hub import InferenceClient
client = InferenceClient("mistralai/Mixtral-8x7B-Instruct-v0.1")
#https://github.com/mshumer/gpt-author/blob/main/Claude_Author.ipynb
#!pip install EbookLib
ANTHROPIC_API_KEY = "YOUR KEY HERE"
stability_api_key = "YOUR KEY HERE" # get it at https://beta.dreamstudio.ai/
def remove_first_line(test_string):
print("removing first line")
if test_string.startswith("Here") and test_string.split("\n")[0].strip().endswith(":"):
return re.sub(r'^.*\n', '', test_string, count=1)
return test_string
def format_prompt(message, history):
prompt = "<s>"
for user_prompt, bot_response in history:
prompt += f"[INST] {user_prompt} [/INST]"
prompt += f" {bot_response}</s> "
prompt += f"[INST] {message} [/INST]"
return prompt
def generate_text(prompt,max_tokens=2000,history=""):
generate_kwargs = dict(
temperature=0.9,
max_new_tokens=max_tokens,
top_p=0.95,
repetition_penalty=1.0,
do_sample=True,
seed=random.randint(1,1000000000)
,
)
#content = prompt_template.format(**prompt_kwargs)
#if VERBOSE:
#print(LOG_PROMPT.format(content))
system_prompt="You are a world-class author. Write the requested content with great skill and attention to detail."
formatted_prompt = format_prompt(f"{system_prompt}, {prompt}", history)
#formatted_prompt = format_prompt(f'{content}', history)
stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
resp = ""
for response in stream:
resp += response.token.text
#yield resp
return resp
def generate_text_og(prompt, model="claude-3-haiku-20240307", max_tokens=2000, temperature=0.7):
headers = {
"x-api-key": ANTHROPIC_API_KEY,
"anthropic-version": "2023-06-01",
"content-type": "application/json"
}
data = {
"model": model,
"max_tokens": max_tokens,
"temperature": temperature,
"system": "You are a world-class author. Write the requested content with great skill and attention to detail.",
"messages": [{"role": "user", "content": prompt}],
}
response = requests.post("https://api.anthropic.com/v1/messages", headers=headers, json=data)
response_text = response.json()['content'][0]['text']
return response_text.strip()
def generate_cover_prompt(plot):
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.")
return response
def generate_title(plot):
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.")
#return remove_first_line(response)
return response
def create_cover_image(plot):
plot = str(generate_cover_prompt(plot))
with open(f"cover.png", "wb") as f: # replace this if running locally, to where you store the cover file
f.write(base64.b64decode(image["base64"]))
def create_cover_image_OG(plot):
plot = str(generate_cover_prompt(plot))
engine_id = "stable-diffusion-xl-beta-v2-2-2"
api_host = os.getenv('API_HOST', 'https://api.stability.ai')
api_key = stability_api_key
if api_key is None:
raise Exception("Missing Stability API key.")
response = requests.post(
f"{api_host}/v1/generation/{engine_id}/text-to-image",
headers={
"Content-Type": "application/json",
"Accept": "application/json",
"Authorization": f"Bearer {api_key}"
},
json={
"text_prompts": [
{
"text": plot
}
],
"cfg_scale": 7,
"clip_guidance_preset": "FAST_BLUE",
"height": 768,
"width": 512,
"samples": 1,
"steps": 30,
},
)
if response.status_code != 200:
raise Exception("Non-200 response: " + str(response.text))
data = response.json()
for i, image in enumerate(data["artifacts"]):
with open(f"/content/cover.png", "wb") as f: # replace this if running locally, to where you store the cover file
f.write(base64.b64decode(image["base64"]))
def generate_chapter_title(chapter_content):
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.")
return remove_first_line(response)
def create_epub(title, author, chapters, cover_image_path='cover.png'):
book = epub.EpubBook()
# Set metadata
book.set_identifier('id123456')
book.set_title(title)
book.set_language('en')
book.add_author(author)
# Add cover image
with open(cover_image_path, 'rb') as cover_file:
cover_image = cover_file.read()
book.set_cover('cover.png', cover_image)
# Create chapters and add them to the book
epub_chapters = []
for i, chapter_content in enumerate(chapters):
chapter_title = generate_chapter_title(chapter_content)
chapter_file_name = f'chapter_{i+1}.xhtml'
epub_chapter = epub.EpubHtml(title=chapter_title, file_name=chapter_file_name, lang='en')
# Add paragraph breaks
formatted_content = ''.join(f'{paragraph.strip()}' for paragraph in chapter_content.split('\n') if paragraph.strip())
epub_chapter.content = f'{chapter_title}{formatted_content}'
book.add_item(epub_chapter)
epub_chapters.append(epub_chapter)
# Define Table of Contents
book.toc = (epub_chapters)
# Add default NCX and Nav files
book.add_item(epub.EpubNcx())
book.add_item(epub.EpubNav())
# Define CSS style
style = '''
@namespace epub "http://www.idpf.org/2007/ops";
body {
font-family: Cambria, Liberation Serif, serif;
}
h1 {
text-align: left;
text-transform: uppercase;
font-weight: 200;
}
'''
# Add CSS file
nav_css = epub.EpubItem(uid="style_nav", file_name="style/nav.css", media_type="text/css", content=style)
book.add_item(nav_css)
# Create spine
book.spine = ['nav'] + epub_chapters
# Save the EPUB file
epub.write_epub(f'{title}.epub', book)
def generate_book(writing_style, book_description, num_chapters):
print("Generating plot outline...")
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."
plot_outline = generate_text(plot_prompt)
print("Plot outline generated.")
chapters = []
for i in range(num_chapters):
print(f"Generating chapter {i+1}...")
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."
chapter = generate_text(chapter_prompt, max_tokens=4000)
chapters.append(remove_first_line(chapter))
print(f"Chapter {i+1} generated.")
time.sleep(1) # Add a short delay to avoid hitting rate limits
print("Compiling the book...")
book = "\n\n".join(chapters)
print("Book generated!")
return plot_outline, book, chapters
def main(writing_style, book_description, num_chapters):
try:
# Generate the book
plot_outline, book, chapters = generate_book(writing_style, book_description, num_chapters)
title = generate_title(plot_outline)
# Save the book to a file
with open(f"{title}.txt", "w") as file:
file.write(book)
#create_cover_image(plot_outline)
# Create the EPUB file
#create_epub(title, 'AI', chapters, 'cover.png')
print(f"Book saved as '{title}.txt'.")
return f'{title}.txt',[f'{title}.txt']
except Exception as e:
return e,None
with gr.Blocks() as app:
with gr.Row():
with gr.Column():
writing_style=gr.Textbox(label="Enter the desired writing style: ")
book_description = gr.Textbox(label="Enter a high-level description of the book: ")
num_chapters = gr.Number(label="Enter the number of chapters: ",precision=0)
with gr.Column():
gen_btn=gr.Button("Generate Book")
outp=gr.HTML()
outf=gr.Files()
gen_btn.click(main,[writing_style,book_description,num_chapters],[outp,outf])
app.queue(default_concurrency_limit=10).launch()