iLL-Ai-Book_writer / book_generator.py
ILLERRAPS's picture
Upload folder using huggingface_hub
4659b69 verified
raw
history blame
4.34 kB
import os
from langchain_groq import ChatGroq
from langchain.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from typing import Dict
import gradio as gr # Import Gradio
import shutil # Import shutil for file operations
# Step 1: Set the environment variable for the Groq API Key
os.environ["GROQ_API_KEY"] = "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" # Use a valid Groq API key
# Step 2: Define a function to create agents
def create_agent(role: str, model_name: str = "llama3-70b-8192", temperature: float = 0.7) -> ChatGroq:
"""Create a LangChain agent for a specific role in book writing."""
prompt_template = ChatPromptTemplate.from_messages([
("system", f"You are a {role}. Write high-quality, engaging content."),
("human", "{input}")
])
llm = ChatGroq(model=model_name, temperature=temperature)
chain = prompt_template | llm | StrOutputParser()
return chain
# Step 3: Create specific agents
chapter_agent = create_agent("chapter writer")
outline_agent = create_agent("outline creator")
editing_agent = create_agent("editor")
proofreading_agent = create_agent("proofreader")
character_agent = create_agent("character developer")
# Step 4: Define functions for generating content
def generate_chapter(title: str, synopsis: str, agent) -> str:
"""Generate a chapter based on the title and synopsis."""
query = f"Write a detailed chapter based on the following synopsis:\n\nTitle: {title}\n\nSynopsis: {synopsis}"
return agent.invoke({"input": query})
def generate_outline(book_title: str, agent) -> Dict[str, str]:
"""Generate an outline for the book."""
query = f"Create an outline for a book titled '{book_title}'."
return agent.invoke({"input": query})
def edit_content(content: str, agent) -> str:
"""Edit the provided content."""
query = f"Edit the following content for clarity and style:\n\n{content}"
return agent.invoke({"input": query})
def proofread_content(content: str, agent) -> str:
"""Proofread the provided content for grammar and spelling errors."""
query = f"Proofread the following content:\n\n{content}"
return agent.invoke({"input": query})
def generate_character_profile(character_name: str, agent) -> str:
"""Generate a character profile for the given character name."""
query = f"Create a character profile for '{character_name}'."
return agent.invoke({"input": query})
def save_chapter_to_file(chapter_text: str) -> str:
"""Save the generated chapter to a .txt file."""
file_path = "generated_chapter.txt"
with open(file_path, "w") as file:
file.write(chapter_text)
return file_path # Return the file path for downloading
# Step 5: Define Gradio interface
def gradio_interface(book_title: str, chapter_title: str, chapter_synopsis: str):
"""Gradio interface for generating book content."""
outline = generate_outline(book_title, outline_agent)
chapter_text = generate_chapter(chapter_title, chapter_synopsis, chapter_agent)
edited_chapter = edit_content(chapter_text, editing_agent)
proofread_chapter = proofread_content(edited_chapter, proofreading_agent)
# Save and download buttons
save_button = gr.Button("Save Chapter")
download_button = gr.Button("Download Chapter")
# Define button actions
save_button.click(fn=save_chapter_to_file, inputs=proofread_chapter, outputs="text")
download_button.click(fn=shutil.copy, inputs="generated_chapter.txt", outputs="text")
return outline, chapter_text, edited_chapter, proofread_chapter
# Step 6: Launch Gradio app
if __name__ == "__main__":
# Define the Gradio interface
iface = gr.Interface(
fn=gradio_interface,
inputs=[
gr.Textbox(label="Book Title"),
gr.Textbox(label="Chapter Title"),
gr.Textbox(label="Chapter Synopsis")
],
outputs=[
gr.Textbox(label="Generated Outline", lines=5),
gr.Textbox(label="Generated Chapter", lines=5),
gr.Textbox(label="Edited Chapter", lines=5),
gr.Textbox(label="Proofread Chapter", lines=5)
],
title="Book Generator",
description="Generate book content using AI."
)
iface.launch(share=True) # Launch the Gradio app