97 lines
2.7 KiB
Python
97 lines
2.7 KiB
Python
|
# %%
|
||
|
|
from typing import Optional, List
|
||
|
|
from dataclasses import field, dataclass
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass
|
||
|
|
class OpusV1Turn:
|
||
|
|
role: str
|
||
|
|
content: str
|
||
|
|
names: List[str] = field(default_factory=list)
|
||
|
|
# If set to true, will not append <|im_end|>, so the model will continue the turn.
|
||
|
|
# In RP you can for example use the following to force a specific character response:
|
||
|
|
# role="text"
|
||
|
|
# names=["Jack"]
|
||
|
|
# open="true"
|
||
|
|
open: bool = False
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass
|
||
|
|
class OpusV1Character:
|
||
|
|
name: str
|
||
|
|
description: str
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass
|
||
|
|
class OpusV1StorySystemPrompt:
|
||
|
|
format: str = "prose"
|
||
|
|
plot_description: str = ""
|
||
|
|
style_description: str = ""
|
||
|
|
characters: List[OpusV1Character] = field(default_factory=list)
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass
|
||
|
|
class OpusV1Prompt:
|
||
|
|
story: Optional[OpusV1StorySystemPrompt] = None
|
||
|
|
turns: List[OpusV1Turn] = field(default_factory=list)
|
||
|
|
|
||
|
|
|
||
|
|
def format_opus_v1_prompt(prompt) -> str:
|
||
|
|
turns = prompt.turns
|
||
|
|
if prompt.story is not None:
|
||
|
|
system = format_opus_v1_system_prompt(prompt.story)
|
||
|
|
turns = [OpusV1Turn(role="system", content=system)] + turns
|
||
|
|
|
||
|
|
parts = []
|
||
|
|
for i, turn in enumerate(turns):
|
||
|
|
assert turn.role in ["user", "text", "system", "assistant"]
|
||
|
|
assert turn.role != "system" or i == 0
|
||
|
|
|
||
|
|
is_last = i == len(turns) - 1
|
||
|
|
open = is_last and turn.open
|
||
|
|
parts.append(format_turn(turn.role, turn.content, turn.names, open=open))
|
||
|
|
return "".join(parts)
|
||
|
|
|
||
|
|
|
||
|
|
def format_turn(
|
||
|
|
role: str, content: str, names: List[str] = [], open: bool = False
|
||
|
|
) -> str:
|
||
|
|
im_start = "<|im_start|>"
|
||
|
|
im_end = "<|im_end|>"
|
||
|
|
|
||
|
|
body = im_start + role
|
||
|
|
if len(names) > 0:
|
||
|
|
body += f" names= {'; '.join(names)}"
|
||
|
|
|
||
|
|
body += "\n"
|
||
|
|
if open:
|
||
|
|
return body + content.lstrip()
|
||
|
|
else:
|
||
|
|
return body + content.strip() + im_end + "\n"
|
||
|
|
|
||
|
|
|
||
|
|
def format_opus_v1_system_prompt(prompt) -> str:
|
||
|
|
format_text = "story" if prompt.format == "prose" else "role-play"
|
||
|
|
system = f"""
|
||
|
|
You are an intelligent, skilled, versatile writer.
|
||
|
|
|
||
|
|
Your task is to write a {format_text} based on the information below.
|
||
|
|
|
||
|
|
Write the {format_text} as if it's a book.
|
||
|
|
""".strip()
|
||
|
|
|
||
|
|
if len(prompt.plot_description) > 0:
|
||
|
|
system += "\n\n\n## Plot description:\n\n"
|
||
|
|
system += prompt.plot_description.strip()
|
||
|
|
if len(prompt.style_description) > 0:
|
||
|
|
system += "\n\n\n## Style description:\n\n"
|
||
|
|
system += prompt.style_description.strip()
|
||
|
|
if len(prompt.characters) > 0:
|
||
|
|
system += "\n\n\n## Characters:\n\n"
|
||
|
|
for character in prompt.characters:
|
||
|
|
system += f"### {character.name}\n\n"
|
||
|
|
system += character.description.strip()
|
||
|
|
system += "\n\n"
|
||
|
|
|
||
|
|
return system.strip()
|