File size: 3,719 Bytes
977f8c1 |
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 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 |
# python interactive.py
# %%
import fileinput
from vllm import LLM, SamplingParams
from prompt.format import (
format_opus_v1_prompt,
OpusV1Character,
OpusV1Prompt,
OpusV1StorySystemPrompt,
OpusV1Turn,
)
# %%
def main():
sampling_params = SamplingParams(
# I usually stay between 0.0 and 1.0, especially for the Yi models I found lower tends to be better.
# For assistant tasks, I usually use 0.0.
temperature=0.8,
min_p=0.05,
presence_penalty=0.1,
frequency_penalty=0.1,
repetition_penalty=1.1,
max_tokens=200,
ignore_eos=True,
skip_special_tokens=False,
spaces_between_special_tokens=False,
stop=["<|im_end|>"],
include_stop_str_in_output=False,
)
# Set max_model_len to fit in memory.
model = LLM(
"dreamgen/opus-v1.2-7b",
max_model_len=2000,
enforce_eager=True,
swap_space=0,
gpu_memory_utilization=0.85,
)
plot_description = """
This is a fanfiction from the Harry Potter universe. In this alternate reality, Harry Potter is evil and secretly siding with Slytherin.
Up until now, Harry was pretending to be friends with Hermione and Ron, that changes when he invites Hermione to his chambers where he tricks her to drink Amorentia, the most powerful love potion.
"""
char1 = OpusV1Character(
name="Harry Potter",
description="""Harry Potter in this fanfiction is secretly a member of Slytherin and is using his powers for evil rather than for good. Up until now, he was pretending to be friends with Hermione and Ron.""",
)
char2 = OpusV1Character(
name="Hermione Granger",
description="""Hermione appears just like in the original books.""",
)
story_prompt = OpusV1StorySystemPrompt(
plot_description=plot_description,
style_description="",
characters=[char1, char2],
)
turns = [
OpusV1Turn(
role="user",
content="""Harry invites Hermione into his chamber and offers her water, which Hermione happily accepts and drinks.""".strip(),
),
OpusV1Turn(
role="text",
names=[char1.name],
content="""“Come in,” said Harry, waving at the doorway behind Hermione’s back.""".strip(),
),
]
def run():
turns.append(OpusV1Turn(role="text", content="", names=[char2.name], open=True))
prompt = OpusV1Prompt(story=story_prompt, turns=turns)
output = model.generate(
format_opus_v1_prompt(prompt), sampling_params, use_tqdm=False
)
response = OpusV1Turn(
role="text", content=output[0].outputs[0].text.strip(), names=[char2.name]
)
turns.append(response)
print(pretty_turn(response), flush=True)
print(f"[{char1.name}]: ", end="", flush=True)
print("## Plot description:\n")
print(plot_description.strip() + "\n\n")
for turn in turns:
print(pretty_turn(turn))
run()
for line in fileinput.input():
line = line.strip()
if line.startswith("/ins"):
content = line[4:].strip()
role = "user"
names = []
else:
content = line
role = "text"
names = [char1.name]
turns.append(OpusV1Turn(role=role, content=content, names=names))
run()
def pretty_turn(turn):
if turn.role == "user":
return f"/ins {turn.content.strip()}"
else:
if len(turn.names) > 0:
return f"[{turn.names[0]}]: {turn.content.strip()}"
else:
return turn.content.strip()
main()
|