model = AutoModelForCausalLM.from_pretrained(
    "ChunB1/Phi-3-interact",
    torch_dtype="auto",
    trust_remote_code=True,
    attn_implementation="flash_attention_2",
)
model.to("cuda")
tokenizer = AutoTokenizer.from_pretrained("microsoft/Phi-3-small-8k-instruct")

pipe = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
)

# Using passage (two sentences), char1 and char2 to classify the interaction type (No, Associating, Thinking, Touching, Observing, Communicating) for one datapoint.
example = {'book_name': '2013_O_Connell,Carol_ItHappensintheDark_MY',
            'sentence_ID': 371,
            'passage': 'Her smile was just a flash, a taste of things to come. He shot her a glance to beg, Play nice.',
            'char1_COREF': 95,
            'char2_COREF': 448,
            'char1': 'He',
            'char2': 'her',}

prompt_five_class_explained = """Communicating: char1 and char2 are engaged in some form of communication, such as speaking, writing, or signaling.
Associating: char1 and char2 are linked by a social or relational context, such as friendship, teamwork, or other associative bonds.
Observing: at least one character is observing or watching another one, without direct interaction.
Thinking: at least one character is thinking about or recalling memories of another one, without direct interaction.
Touching: char1 and char2 are engaged in physical touch or contact."""

prompt_base = "what kind of interaction between char1 and char2? Choose one of six options: No, Associating, Thinking, Touching, Observing, Communicating."

prompt = """Task Description: Classify the type of interaction between char1 and char2 in a given passage. There are six categories of interaction:

No interaction: Direct or indirect interaction does not occur between char1 and char2. Any imagination or assumption of interaction also counts as No.
""" + prompt_five_class_explained + prompt_base

prompt_suffix = "Only return the option and don't provide any extra information."

prompt_full = f"passage: {example['passage']}, char1: {example['char1']}, char2: {example['char2']}, " + prompt + prompt_suffix
messages = [{"role": "user", "content": prompt_full}]

generation_args = {
    "max_new_tokens": 15,
    "return_full_text": False,
    "do_sample": False,
}

# Label will be "Communicating"
print(pipe(messages, **generation_args)[0]["generated_text"].strip())