ExternalVendorCode/lollms
1
0
Fork 0
mirror of https://github.com/ParisNeo/lollms.git synced 2024-12-19 04:37:54 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Added illustration generation

Browse Source
This commit is contained in:
Saifeddine ALOUI 2024-06-02 01:19:32 +02:00
parent 9a94454661
commit 70f73d622e
2 changed files with 75 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

19
lollms/functions/markdown2latex.py
View File

@ -28,7 +28,8 @@ def markdown_to_latex(file_path: str) -> str:
# Define conversion rules from markdown to LaTeX
conversion_rules = [
(r'\\', r'\\textbackslash{}'), # Escape backslashes
(r'([#]+) (.*)', lambda m: '\\' + 'sub'*(len(m.group(1))-1) + 'section{' + m.group(2) + '}'), # Convert headings
(r'^## (Chapter\s*\d+:\s*)?(.*)', lambda match: f"\\chapter{{{match.group(2)}}}"), # Convert chapters
(r'^### (.*)', r'\\section*{\1}'), # Convert subsections without numbering
(r'\*\*(.*?)\*\*', r'\\textbf{\1}'), # Bold text
(r'\*(.*?)\*', r'\\textit{\1}'), # Italic text
(r'\!\[(.*?)\]\((.*?)\)', r'\\begin{figure}[h!]\n\\centering\n\\includegraphics[width=\\textwidth]{\2}\n\\caption{\1}\n\\end{figure}'), # Images
@ -43,14 +44,26 @@ def markdown_to_latex(file_path: str) -> str:
# Apply conversion rules
latex_text = markdown_text
for pattern, replacement in conversion_rules:
latex_text = re.sub(pattern, replacement, latex_text, flags=re.MULTILINE)
if callable(replacement):
latex_text = re.sub(pattern, replacement, latex_text, flags=re.MULTILINE)
else:
latex_text = re.sub(pattern, replacement, latex_text, flags=re.MULTILINE)
# Create the LaTeX document structure
# Add lettrine to the first letter of the first paragraph of each chapter
latex_text = re.sub(r'(\\chapter{[^}]*}\n\n?)(\w)', r'\1\\lettrine[lines=2]{\2}{}', latex_text)
# Create the LaTeX document structure with fancyhdr for headers
latex_document = f"""
\\documentclass{{book}}
\\usepackage{{hyperref}}
\\usepackage{{graphicx}}
\\usepackage{{verbatim}}
\\usepackage{{fancyhdr}}
\\usepackage{{lettrine}}
\\pagestyle{{fancy}}
\\fancyhf{{}}
\\fancyhead[LE,RO]{{\\thepage}}
\\fancyhead[RE,LO]{{\\leftmark}}
\\begin{{document}}

66
lollms/functions/story_writing.py
View File

@ -20,11 +20,53 @@ import json
# Import markdown2latex
from lollms.functions.markdown2latex import markdown_to_latex
from lollms.functions.file_manipulation import change_file_extension
from lollms.functions.generate_image import build_image
from lollms.client_session import Client
import copy
def needs_illustration(section_name, mydict):
for section in mydict['sections']:
if section['section_name'] == section_name:
return section['build_illustration']
return False # Return None if the section is not found
def build_section_illustration(llm, prompt_ideas, current_section, content, client:Client):
discussion_prompt_separator = llm.config.discussion_prompt_separator
start_header_id_template = llm.config.start_header_id_template
end_header_id_template = llm.config.end_header_id_template
separator_template = llm.config.separator_template
system_message_template = llm.config.system_message_template
prompt = "\n".join([
f"{start_header_id_template}{system_message_template}{end_header_id_template}",
f"You are illustration builder for a book section. Your task is to build a prompt to generate an illustration from a book section of a story.",
"Your output must strictly contain only a json formatted output like this:",
"```json",
"{",
'"prompt":"image generation prompt",'
'"width":the width of the image,'
'"height":the height of the image'
"}"
"```",
f"{start_header_id_template}General idea{end_header_id_template}",
f"{prompt_ideas}\n\n"
f"{start_header_id_template}Section title{end_header_id_template}",
f"{current_section}",
f"{start_header_id_template}Section content{end_header_id_template}",
f"{content}",
f"{start_header_id_template}story_section_illustrator{end_header_id_template}"
]
)
image_generation_prompt = llm.fast_gen(prompt, callback=llm.sink)
code_blocks = llm.extract_code_blocks(image_generation_prompt)
if len(code_blocks)>0:
code = json.loads(code_blocks[0]["content"])
return build_image(image_generation_prompt,width=code["width"],height=code["height"], client=client)
# Define the core functions
def start_writing_story(prompt_ideas: str, llm: Any, story_file_path: str, build_latex:bool=False) -> str:
def start_writing_story(prompt_ideas: str, llm: Any, story_file_path: str, build_latex:bool=False, client:Client = None) -> str:
discussion_prompt_separator = llm.config.discussion_prompt_separator
start_header_id_template = llm.config.start_header_id_template
end_header_id_template = llm.config.end_header_id_template
@ -44,6 +86,7 @@ def start_writing_story(prompt_ideas: str, llm: Any, story_file_path: str, build
"{\n"
"\"section_name\": \"The name of the section\",\n"
"\"section_description\": \"A short description of the section\"\n"
"\"build_illustration\": A boolean indicating if an illustration is needed for this section\n"
"}, ...\n"
"]\n"
"}\n"
@ -79,7 +122,9 @@ def start_writing_story(prompt_ideas: str, llm: Any, story_file_path: str, build
story_file_path=story_file_path,
story_plan=story_plan,
current_section=section_name,
prompt_ideas=prompt_ideas
prompt_ideas=prompt_ideas,
add_illustration=needs_illustration(section_name, story_plan),
client=client
)
section_full["content"]=new_section
final_story_content += f"\n## {section_name}\n\n{new_section}\n"
@ -116,10 +161,10 @@ def start_writing_story(prompt_ideas: str, llm: Any, story_file_path: str, build
return trace_exception(e)
# Define the metadata functions
def start_writing_story_function(llm, story_file_path, build_latex:bool=False):
def start_writing_story_function(llm, story_file_path, build_latex:bool=False, client:Client=None):
return {
"function_name": "start_writing_story",
"function": partial(start_writing_story, llm=llm, story_file_path=story_file_path, build_latex=build_latex),
"function": partial(start_writing_story, llm=llm, story_file_path=story_file_path, build_latex=build_latex, client=client),
"function_description": "Starts writing a story based on the provided prompt ideas, generating a plan in JSON format, and writing the story section by section.",
"function_parameters": [
{"name": "prompt_ideas", "type": "str"}
@ -128,7 +173,7 @@ def start_writing_story_function(llm, story_file_path, build_latex:bool=False):
# Define the core function
def write_story_section(prompt_ideas: str, llm: Any, story_file_path: str, story_plan: dict, current_section: str) -> str:
def write_story_section(prompt_ideas: str, llm: Any, story_file_path: str, story_plan: dict, current_section: str, add_illustration:bool=False, client:Client=None) -> str:
discussion_prompt_separator = llm.config.discussion_prompt_separator
start_header_id_template = llm.config.start_header_id_template
end_header_id_template = llm.config.end_header_id_template
@ -162,8 +207,11 @@ def write_story_section(prompt_ideas: str, llm: Any, story_file_path: str, story
f"{start_header_id_template}story_section_writer{end_header_id_template}"
]
)
new_section += llm.fast_gen(prompt, callback=llm.sink)
content = llm.fast_gen(prompt, callback=llm.sink)
new_section += content
if add_illustration and client:
illustration = build_section_illustration(llm, prompt_ideas, current_section, content, client)
new_section += "\n" + illustration
# Write the new section to the story file
story_path.write_text(new_section)
@ -194,6 +242,10 @@ def write_story_section(prompt_ideas: str, llm: Any, story_file_path: str, story
new_section = f"## {current_section}\n\n"
new_section += llm.fast_gen(prompt, callback=llm.sink).strip()
if add_illustration and client:
illustration = build_section_illustration(llm, prompt_ideas, current_section, content, client)
new_section += "\n" + illustration
# Append the new section to the story file
story_path.write_text(story_content + "\n" + new_section)