Generating TV shows with template chains

We can use template chains to generate text that follows a specific format. In this example, we will generate the details of a fictional TV show using a mix of accurate and creative inference arguments.

Refer to the text generation example for information on how to load models. This example assumes that the model name "model" is registered with an instruct model compatible with the Alpaca format (i.e. WizardLM).

Setting the inference arguments

accurate_infer_args = InferArgs(top_p=0.3)
creative_infer_args = InferArgs(
  min_new_tokens=200,
  typical_mass=0.8,
  rep_penalty=1.1,
)

The accurate_infer_args object specifies that we want to use top-p sampling for accuracy, while the creative_infer_args object is used to specify the minimum number of new tokens, typical mass, and replication penalty for the language model to generate creative and novel ideas.

Setting up our instruction chain

We want our language model to generate the details by filling in the brackets in this format:

### Instruction
Write the details and a brief summary for a completely new television show about the topic below:

(insert TV show idea here)

### Response
Title: {title}
Genre: {genre}
Production studio: {prod_studio}
Number of episodes: {no_episodes}
Plot: {plot}

We create the instruct_chain object, where we can define the chain, its nodes, and the type of its argument and return keys.

instruct_chain = ld.chain(
  args={"instruction": str},
  returns={
    "title": str,
    "genre": str,
    "no_episodes": int,
    "prod_studio": str,
    "summary": str,
  },
  nodes=[
    ld.format_args(
      """\
### Instruction
Write the details and a brief summary for a completely new television show about the topic below:

{instruction}

### Response
"""
    ),
    "Title: ",
    ld.returns("title", "\n", infer_args=accurate_infer_args),
    "Genre: ",
    ld.returns("genre", "\n", infer_args=accurate_infer_args),
    "Production studio: ",
    ld.returns("prod_studio", "\n", infer_args=accurate_infer_args),
    "Number of episodes: ",
    ld.returns("no_episodes", "\n", infer_args=accurate_infer_args),
    "Plot: ",
    ld.returns("summary", "\n", infer_args=creative_infer_args),
  ],
)

The Title:, Genre:, Production studio:, Number of episodes:, and Plot: strings are followed by the ld.returns function to retrieve the corresponding outputs from the language model. The second argument of the ld.returns function specifies the end of each generation, here it is marked by a newline character. The infer_args parameter is passed to the function to specify the inference arguments for each node.

Streaming the response

We can then stream the response. The code below streams the output of a chain by iterating through each response generated, printing the corresponding key and value pairs accordingly.

for resp in instruct_chain.stream(
  args={"instruction": instruction}
  ):
  if isinstance(resp, RespReturns):
    print("\n")
    print(resp.key, end=": ", flush=True)
  elif isinstance(resp, RespInfer):
    print(resp.tokstr, end="", flush=True)

The responses yielded by the stream() method will follow this format:

RespReturns(key="return key")
RespInfer(tokid=1234, tokstr="token string", running_infer="string containing everything generated so far")
...

If the response yielded is of type RespReturns, it prints the header. If the response is of type RespInfer, it prints the value generated by the language model, token-by-token.

Sample output

title: Shadow World

genre: Fantasy Anime

prod_studio: XYZ Animation

no_episodes: 26

summary: Shadow World is an epic fantasy adventure that takes place in a world where the line between good and evil is constantly blurred. The story follows a young hero named Tsubasa, who discovers that he was born with the power to control shadows – dark, mystical creatures that serve as guardians of both good and evil.

Getting output as a dictionary

Alternatively, we can get the generation output as a dictionary through the call() method:

print(instruct_chain.call(args={"instruction": instruction}))

Sample output

For the input A fantasy anime:

LDResult(returns={'title': 'Shadow World', 'genre': 'Fantasy Anime', 'prod_studio': 'XYZ Animation', 'no_episodes': 26, 'summary': 'Shadow World is an epic fantasy adventure that takes place in a world where the line between good and evil is constantly blurred. The story follows a young hero named Tsubasa, who discovers that he was born with the power to control shadows – dark, mystical creatures that serve as guardians of both good and evil.'}, prompt_tokens=35, completion_tokens=88)