File size: 2,726 Bytes
814f4ef |
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 92 93 94 95 96 97 |
# %%
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()
|