Upload /example/prompt/format.py with huggingface_hub
Browse files- example/prompt/format.py +96 -0
example/prompt/format.py
ADDED
@@ -0,0 +1,96 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# %%
|
2 |
+
from typing import Optional, List
|
3 |
+
from dataclasses import field, dataclass
|
4 |
+
|
5 |
+
|
6 |
+
@dataclass
|
7 |
+
class OpusV1Turn:
|
8 |
+
role: str
|
9 |
+
content: str
|
10 |
+
names: List[str] = field(default_factory=list)
|
11 |
+
# If set to true, will not append <|im_end|>, so the model will continue the turn.
|
12 |
+
# In RP you can for example use the following to force a specific character response:
|
13 |
+
# role="text"
|
14 |
+
# names=["Jack"]
|
15 |
+
# open="true"
|
16 |
+
open: bool = False
|
17 |
+
|
18 |
+
|
19 |
+
@dataclass
|
20 |
+
class OpusV1Character:
|
21 |
+
name: str
|
22 |
+
description: str
|
23 |
+
|
24 |
+
|
25 |
+
@dataclass
|
26 |
+
class OpusV1StorySystemPrompt:
|
27 |
+
format: str = "prose"
|
28 |
+
plot_description: str = ""
|
29 |
+
style_description: str = ""
|
30 |
+
characters: List[OpusV1Character] = field(default_factory=list)
|
31 |
+
|
32 |
+
|
33 |
+
@dataclass
|
34 |
+
class OpusV1Prompt:
|
35 |
+
story: Optional[OpusV1StorySystemPrompt] = None
|
36 |
+
turns: List[OpusV1Turn] = field(default_factory=list)
|
37 |
+
|
38 |
+
|
39 |
+
def format_opus_v1_prompt(prompt) -> str:
|
40 |
+
turns = prompt.turns
|
41 |
+
if prompt.story is not None:
|
42 |
+
system = format_opus_v1_system_prompt(prompt.story)
|
43 |
+
turns = [OpusV1Turn(role="system", content=system)] + turns
|
44 |
+
|
45 |
+
parts = []
|
46 |
+
for i, turn in enumerate(turns):
|
47 |
+
assert turn.role in ["user", "text", "system", "assistant"]
|
48 |
+
assert turn.role != "system" or i == 0
|
49 |
+
|
50 |
+
is_last = i == len(turns) - 1
|
51 |
+
open = is_last and turn.open
|
52 |
+
parts.append(format_turn(turn.role, turn.content, turn.names, open=open))
|
53 |
+
return "".join(parts)
|
54 |
+
|
55 |
+
|
56 |
+
def format_turn(
|
57 |
+
role: str, content: str, names: List[str] = [], open: bool = False
|
58 |
+
) -> str:
|
59 |
+
im_start = "<|im_start|>"
|
60 |
+
im_end = "<|im_end|>"
|
61 |
+
|
62 |
+
body = im_start + role
|
63 |
+
if len(names) > 0:
|
64 |
+
body += f" names= {'; '.join(names)}"
|
65 |
+
|
66 |
+
body += "\n"
|
67 |
+
if open:
|
68 |
+
return body + content.lstrip()
|
69 |
+
else:
|
70 |
+
return body + content.strip() + im_end + "\n"
|
71 |
+
|
72 |
+
|
73 |
+
def format_opus_v1_system_prompt(prompt) -> str:
|
74 |
+
format_text = "story" if prompt.format == "prose" else "role-play"
|
75 |
+
system = f"""
|
76 |
+
You are an intelligent, skilled, versatile writer.
|
77 |
+
|
78 |
+
Your task is to write a {format_text} based on the information below.
|
79 |
+
|
80 |
+
Write the {format_text} as if it's a book.
|
81 |
+
""".strip()
|
82 |
+
|
83 |
+
if len(prompt.plot_description) > 0:
|
84 |
+
system += "\n\n\n## Plot description:\n\n"
|
85 |
+
system += prompt.plot_description.strip()
|
86 |
+
if len(prompt.style_description) > 0:
|
87 |
+
system += "\n\n\n## Style description:\n\n"
|
88 |
+
system += prompt.style_description.strip()
|
89 |
+
if len(prompt.characters) > 0:
|
90 |
+
system += "\n\n\n## Characters:\n\n"
|
91 |
+
for character in prompt.characters:
|
92 |
+
system += f"### {character.name}\n\n"
|
93 |
+
system += character.description.strip()
|
94 |
+
system += "\n\n"
|
95 |
+
|
96 |
+
return system.strip()
|