Is this correct approach to do Prompt Tuning on DollyV2 model

🤗Transformers
1

I am using this link to study about Prompt Tuning. Parameter-Efficient Fine-Tuning using :hugs: PEFT

Please let me know if its correct way to do Prompt tuning and saving the model.

It has 4 options. I am interested in option 3.

  1. Prompt Tuning: The Power of Scale for Parameter-Efficient Prompt Tuning

Can I use Prompt Tuning to tune DollyV2 [databricks/dolly-v2-12b · Hugging Face] model.

The use-case is:

I have Context which has lot of paragraphs and then Question , the model has to answer the Question based on Context in a professional manner. Also can it classify the Question as relevant if answer is present in Context and irrelevant if answer is not in Context

The code that I have written is:

    task_type=TaskType.CAUSAL_LM,
    prompt_tuning_init=PromptTuningInit.TEXT,
    num_virtual_tokens=8,
    prompt_tuning_init_text="Answer the question as truthfully as possible using and only using the provided context and if the answer is not contained within the context/text, say Irrelevant",
    tokenizer_name_or_path="dolly-v2-3b"
)

tokenizer = AutoTokenizer.from_pretrained("dolly-v2-3b")
model = AutoModelForCausalLM.from_pretrained("dolly-v2-3b",load_in_8bit=True,device_map='auto') #,load_in_8bit=True

model = get_peft_model(model, peft_config)

train_data = [
    {
        "Context": "How to Link Credit Card to ICICI Bank Account Step 1: Login to ICICIBank.com using your existing internet banking credentials. Step 2: Go to the 'Service Request' section. Step 3: Visit the 'Customer Service' option. Step 4: Select the Link Accounts/ Policy option to link your credit card to the existing user ID.",
        "Question": "How to add card?",
        "Answer": "Relevant. To add your card you can follow these steps: Step 1: Login to ICICIBank.com using your existing internet banking credentials. Step 2: Go to the 'Service Request' section. Step 3: Visit the 'Customer Service' option. Step 4: Select the Link Accounts/ Policy option to link your credit card to the existing user ID."
    },
    {
        "Context": "The python programming language is used in many different fields including web development, data analysis, artificial intelligence and scientific computing. It is a high-level language that is easy to learn and has a large community of users who can provide support and advice. ",
        "Question": "What is Python used for?",
        "Answer": "Relevant. Python is used in many different fields including web development, data analysis, artificial intelligence and scientific computing."
    },
    {
        "Context": "The United States is a federal republic consisting of 50 states, a federal district, five major self-governing territories, and various possessions. It has a population of over 330 million people and is the third most populous country in the world. The capital and largest city is Washington, D.C.",
        "Question": "What is the population of the United States?",
        "Answer": "Relevant. The United States has a population of over 330 million people."
    }]

Define a function to map examples to inputs and targets

def preprocess_function(examples):
    tokenized_examples = tokenizer(
        examples["Question"][0],
        examples["Context"][0],
        truncation=True,
        max_length=1024,
        padding="max_length"
    )
    tokenized_examples['labels']=tokenizer(
        examples["Answer"],
        truncation=True,
        max_length=1024,
        padding="max_length",
        return_tensors="pt")['input_ids'][0]
        
    return tokenized_examples

tokenized_train_data = [preprocess_function(example) for example in train_data]


class DemoDataset(Dataset):
    def __init__(self, data):
        self.data = data

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        sample = self.data[idx]
        
        item = {k: torch.tensor(v) for k, v in sample.items()}
        return item

dataset = DemoDataset(tokenized_train_data)

training_args = TrainingArguments(
    output_dir="results",
    learning_rate=1e-5,
    per_device_train_batch_size=1,
    num_train_epochs=10,
    weight_decay=0.01,
    logging_steps=1,
    save_steps=1,
    logging_dir="logs"
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
    # data_collator=data_collator,
    tokenizer=tokenizer
)
trainer.train()

trainer.save_model("dolly3b_demo_model")

Inference

from peft import PeftModel, PeftConfig
tokenizer = AutoTokenizer.from_pretrained("dolly-v2-3b")
model = AutoModelForCausalLM.from_pretrained("dolly3b_demo_model")
model = get_peft_model(model, peft_config)


# Define example
context = "How to Link Credit Card to ICICI Bank Account Step 1: Login to ICICIBank.com using your existing internet banking credentials. Step 2: Go to the 'Service Request' section. Step 3: Visit the 'Customer Service' option. Step 4: Select the Link Accounts/ Policy option to link your credit card to the existing user ID."
question = "How to add card?"

# Encode inputs with prompt and tokenize
inputs = [f"{context} {question}"]
inputs_encoded = tokenizer(inputs, padding=True, truncation=True, max_length=1024, return_tensors="pt")

outputs = model.generate(input_ids=inputs_encoded["input_ids"], attention_mask=inputs_encoded["attention_mask"], max_new_tokens=200, eos_token_id=3)
print(tokenizer.batch_decode(outputs.detach().cpu().numpy(), skip_special_tokens=True))

Is the inference and training code correct. ? The inference take lot of time to generate.

@ pratikbhavsar @valhalla @ Lewis Tunstall @ Merve Noyan @ Lucile Saulnier