|
|
|
from typing import Optional, List |
|
from dataclasses import field, dataclass |
|
|
|
|
|
@dataclass |
|
class OpusV1Turn: |
|
role: str |
|
content: str |
|
names: List[str] = field(default_factory=list) |
|
|
|
|
|
|
|
|
|
|
|
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() |
|
|
