DreamGenX commited on
Commit
814f4ef
1 Parent(s): b0ebbd2

Upload /example/prompt/format.py with huggingface_hub

Browse files
Files changed (1) hide show
  1. 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()