example/interactive.py

# 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()