SHORTFORM_TO_FULL_TASK_TYPES = {
    "exqa": "extractive question answering",
    "mcqa": "multiple-choice question answering",
    "qg": "question generation",
    "qa": "question answering without choices",
    "ynqa": "yes-no question answering",
    "coref": "coreference resolution",
    "paraphrase": "paraphrase generation",
    "paraphrase_id": "paraphrase identification",
    "sent_comp": "sentence completion",
    "sentiment": "sentiment",
    "summarization": "summarization",
    "text_gen": "text generation",
    "topic_class": "topic classification",
    "wsd": "word sense disambiguation",
    "te": "textual entailment",
    "nli": "natural language inference",
}

import torch

# test data
sample_paragraph = """1. “Confidential Information”, whenever used in this Agreement, shall mean any data, document, specification and other information or material, that is delivered or disclosed by UNHCR to the Recipient in any form whatsoever, whether orally, visually in writing or otherwise (including computerized form), and that, at the time of disclosure to the Recipient, is designated as confidential."""

# load model
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model_id = 'mychen76/Llama-3.1-8B-bonito-v1-awq'
bonito = AWQBonito(model_id)

from transformers import set_seed
from pprint import pprint
set_seed(2)

# Generate synthetic instruction tuning dataset
sampling_params = {
    "max_new_tokens": 256,
    "top_p": 0.95,
    "temperature": 0.7,
    "num_return_sequences": 1,
}
synthetic_dataset = bonito.generate_task(
    unannotated_paragraph, task_type="nli", sampling_params=sampling_params
)
pprint("----Generated Instructions----")
pprint(f'Input: {synthetic_dataset["input"]}')
pprint(f'Output: {synthetic_dataset["output"]}')

'----Generated Instructions----'
('Input: Given that 1. “Confidential Information”, whenever used in this '
 'Agreement, shall mean any data, document, specification and other '
 'information or material, that is delivered or disclosed by UNHCR to the '
 'Recipient in any form whatsoever, whether orally, visually in writing or '
 'otherwise (including computerized form), and that, at the time of disclosure '
 'to the Recipient, is designated as confidential. Does it follow that '
 'Confidential Information is not used in this Agreement. Yes, no, or maybe?')
'Output: No'

from transformers import set_seed
from pprint import pprint
set_seed(2)

# Generate synthetic instruction tuning dataset
sampling_params = {
    "max_new_tokens": 256,
    "top_p": 0.95,
    "temperature": 0.7,
    "num_return_sequences": 1,
}
synthetic_dataset = bonito.generate_task(
    unannotated_paragraph, task_type="qa", sampling_params=sampling_params
)
pprint("----Generated Instructions----")
pprint(f'Input: {synthetic_dataset["input"]}')
pprint(f'Output: {synthetic_dataset["output"]}')

'----Generated Instructions----'
('Input: Read the following context and answer the question.\n'
 '\n'
 'Context: 1. “Confidential Information”, whenever used in this Agreement, '
 'shall mean any data, document, specification and other information or '
 'material, that is delivered or disclosed by UNHCR to the Recipient in any '
 'form whatsoever, whether orally, visually in writing or otherwise (including '
 'computerized form), and that, at the time of disclosure to the Recipient, is '
 'designated as confidential.\n'
 'Question: What is confidential information?\n'
 'Answer:')
('Output: information that is designated as confidential at the time of '
 'disclosure')

# test data
from datasets import load_dataset

# load dataaset with unannotated text
unannotated_text_ds = load_dataset(
    "BatsResearch/bonito-experiment",
    "unannotated_contract_nli"
)["train"].select(range(10))

from vllm import LLM, SamplingParams
import torch 

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model_id = 'mychen76/Llama-3.1-8B-bonito-v1-awq'

# Create an LLM.
vbonito = VLLMBonito(model=model_id, 
              quantization="AWQ",
              gpu_memory_utilization=0.80,
              max_seq_len_to_capture=4096,
              max_model_len=4096,
              speculative_max_model_len = 4096)

from vllm import SamplingParams
from transformers import set_seed
from pprint import pprint
set_seed(2)

# Generate synthetic instruction tuning dataset
sampling_params = SamplingParams(max_tokens=256, top_p=0.95, temperature=0.5, n=1)

synthetic_dataset = vbonito.generate_tasks(
    unannotated_text_ds,
    context_col="input",
    task_type="nli",
    sampling_params=sampling_params
)

pprint("----Generated Instructions----")
pprint(f'Input: {synthetic_dataset["input"]}')
pprint(f'Output: {synthetic_dataset["output"]}')

("Input: ['Given that 2.3 Provided that the Recipient has a written agreement "
 'with the following persons or entities requiring them to treat the '
 'Confidential Information in accordance with this Agreement, the Recipient '
 'may disclose the Confidential Information to: 2.3.1  Any other party with '
 'the Discloser’s prior written consent; and 2.3.2 the Recipient’s employees, '
 'officials, representatives and agents who have a strict need to know the '
 'contents of the Confidential Information, and employees, officials, '
 'representatives and agents of any legal entity that it controls, controls '
 'it, or with which it is under common control, who have a similar need to '
 'know the contents of the Confidential Information, provided that, for these '
 'purposes a controlled legal entity means: Does it follow that The Recipient '
 "is a person Yes, no, or maybe?', '5. All Confidential Information in any "
 'form and any medium, including all copies thereof, disclosed to the '
 'Recipient shall be returned to UNHCR or destroyed:  (a) if a business '
 'relationship is not entered into with UNHCR on or before the date which is '
 'three (3) months after the date both Parties have signed the Agreement; or '
 'Based on that information, is the claim: "The UNHCR is not a recipient." '
 "true, false, or inconclusive?', 'Given 4. Nothing in this Agreement is to be "
 'construed as granting the Recipient, by implication or otherwise, any right '
 'whatsoever with respect to the Confidential Information or part thereof. Is '
 'it guaranteed true that "The Recipient is not granted any rights to the '
 'Confidential Information."? Yes, no, or maybe?\', \'11. The Recipient shall '
 'not advertise or otherwise make public the fact that it has a confidential '
 'relationship with UNHCR, nor shall the Recipient, in any manner whatsoever '
....