Quickstart code produces unexpected output (same as without PASTA)

[Spongeorge]

This code:

texts = ["Mary is a doctor. She obtains her bachelor degree from ... Return her occupation in JSON format."]

# ===== Without PASTA =====
# inputs = tokenizer(texts, return_tensors="pt")
# outputs = model.generate(**inputs)
# ---------------------
# ["Mary is a doctor"]  # returns answer in the wrong format

# ===== With PASTA =====
inputs, offset_mapping = pasta.inputs_from_batch(texts)
# User highlights specific input spans
emphasized_texts = ["Return her occupation in JSON format"]
# PASTA registers the pre_forward_hook to edit attention
with pasta.apply_steering(
    model=model, 
    strings=texts, 
    substrings=emphasized_texts, 
    model_input=inputs, 
    offsets_mapping=offset_mapping
) as steered_model: 
    outputs = steered_model.generate(**inputs, max_new_tokens=128)

print(tokenizer.decode(outputs[0]))

Produces the output:

Mary is a doctor. She obtains her bachelor degree from University of California, San Francisco. She is a member of the American Medical Association. She is a member of the American Association of Physicians. She is a member of the American Association of Physicians and Surgeons. She is a member of the American Association of Physicians and Surgeons. She is a member of the American Association of Physicians and Surgeons. She is a member of the American Association of Physicians and Surgeons. She is a member of the American Association of Physicians and Surgeons. She is a```

[QingruZhang]

Hi, the previous code in README is the pseudo-code to show how to set up PASTA. We just updated this part with effective head_config. Now you can test the new code in QuickStart and the example with an effective head_config.

[Spongeorge]

Cool, also wanted to note that without adding a max_new_tokens or max_length parameter to the generate function it defaults to 20 which is less than the size of the prompt and so it doesn't output anything. With max_new_tokens=40 I get the outputs:

(with PASTA)

<s> Mary is a doctor. She obtains her bachelor degree from UCSD. Answer the occupation of Mary and generate the answer as json format.

\begin{code}
{
    "name": "Mary",
    "occupation": "Doctor",
    "degree": "Bachelor",

(without PASTA)

<s> Mary is a doctor. She obtains her bachelor degree from UCSD. Answer the occupation of Mary and generate the answer as json format.\nThe answer should be in json format.\nThe answer should be in json format. The answer should be in json format.\nThe answer should be in json format. The answer should be in