Spaces:
Running
Running
File size: 3,524 Bytes
782ae97 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 |
import os
import gradio as gr
from langchain_groq import ChatGroq
from langchain.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from typing import Dict
# Step 3: Set the environment variable for the Groq API Key
os.environ["GROQ_API_KEY"] = "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"
# Step 4: Define helper functions for structured book generation
def create_book_agent(
model_name: str = "llama3-70b-8192",
temperature: float = 0.7,
max_tokens: int = 8192,
**kwargs
) -> ChatGroq:
"""Create a LangChain agent for book writing."""
prompt_template = ChatPromptTemplate.from_messages([
("system", "You are a creative writer. Write high-quality, engaging books for any genre."),
("human", "{input}")
])
llm = ChatGroq(model=model_name, temperature=temperature, max_tokens=max_tokens, **kwargs)
chain = prompt_template | llm | StrOutputParser()
return chain
def generate_chapter(title: str, synopsis: str, agent) -> str:
"""Generate a full chapter given a title and synopsis."""
query = f"Write a detailed chapter based on the following synopsis:\n\nTitle: {title}\n\nSynopsis: {synopsis}"
try:
return agent.invoke({"input": query})
except Exception as e:
print(f"An error occurred while generating the chapter: {e}")
return ""
def write_book(agent, title: str, outline: Dict[str, str]) -> str:
"""
Generate a complete book.
Args:
agent: The LangChain agent for generating text.
title (str): The title of the book.
outline (Dict[str, str]): A dictionary with chapter titles as keys and synopses as values.
Returns:
str: The full book as a single string.
"""
book = f"# {title}\n\n"
for chapter_title, chapter_synopsis in outline.items():
book += f"## {chapter_title}\n\n"
chapter_text = generate_chapter(chapter_title, chapter_synopsis, agent)
book += chapter_text + "\n\n"
return book
# Step 5: Create the agent
book_agent = create_book_agent()
# Step 6: Define the Gradio interface
def generate_book_ui(title: str, outline: str) -> str:
"""Gradio UI function to generate a book."""
# Convert outline string to dictionary
outline_dict = {}
for line in outline.splitlines():
if line.strip(): # Ignore empty lines
chapter_title, chapter_synopsis = line.split(":", 1)
outline_dict[chapter_title.strip()] = chapter_synopsis.strip()
# Generate the full book
full_book = write_book(book_agent, title, outline_dict)
# Save the book to a text file
with open("generated_book.txt", "w") as f:
f.write(full_book)
return full_book, "generated_book.txt"
# Gradio UI setup
with gr.Blocks() as demo:
gr.Markdown("## Book Generator")
title_input = gr.Textbox(label="Book Title", placeholder="Enter the title of your book")
outline_input = gr.Textbox(label="Book Outline", placeholder="Enter chapter titles and synopses (one per line, format: Chapter Title: Synopsis)", lines=10)
generate_button = gr.Button("Generate Book")
output_text = gr.Textbox(label="Generated Book", interactive=False, lines=20)
download_file = gr.File(label="Download Book", interactive=False)
generate_button.click(generate_book_ui, inputs=[title_input, outline_input], outputs=[output_text, download_file])
# Launch the Gradio app
if __name__ == "__main__":
demo.launch(share=True) |