granite-3.2-8b-alora-requirement-check

pip install git+ssh://git@github.com:IBM/activated-lora.git

import torch,os
from transformers import AutoTokenizer,  AutoModelForCausalLM, DynamicCache
from alora.peft_model_alora import aLoRAPeftModelForCausalLM
from alora.config import aLoraConfig
from alora.tokenize_alora import tokenize_alora

REUSE_CACHE = False  

token = os.getenv("HF_MISTRAL_TOKEN")
BASE_NAME = "ibm-granite/granite-3.2-8b-instruct"
LORA_NAME = "ibm-granite/granite-3.2-8b-alora-requirement-check"
device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# Load model
tokenizer = AutoTokenizer.from_pretrained(BASE_NAME,padding_side='left',trust_remote_code=True, token=token)
model_base = AutoModelForCausalLM.from_pretrained(BASE_NAME,device_map="auto")
model_req_check = aLoRAPeftModelForCausalLM.from_pretrained(model_base, LORA_NAME)

question = "What is IBM Research?"
print("Question:" + question)
question_chat = [
    {
        "role": "user",
        "content": question
    },
]

# Generate answer with base model
input_text = tokenizer.apply_chat_template(question_chat,tokenize=False,add_generation_prompt=True)
# Remove default system prompt
len_sys = len(input_text.split("<|start_of_role|>user"))
input_text = input_text[len_sys:]

#tokenize
inputs = tokenizer(input_text, return_tensors="pt")
if REUSE_CACHE: #save KV cache for future aLoRA call
    prompt_cache = DynamicCache()
    with model_req_check.disable_adapter():
        output_dict = model_base.generate(inputs["input_ids"].to(device), attention_mask=inputs["attention_mask"].to(device), max_new_tokens=600,past_key_values = prompt_cache, return_dict_in_generate=True)
    answer_cache = output_dict.past_key_values
    output = output_dict.sequences
else: #simplest call
    with model_req_check.disable_adapter():
        output = model_req_check.generate(inputs["input_ids"].to(device), attention_mask=inputs["attention_mask"].to(device), max_new_tokens=600)
output_text = tokenizer.decode(output[0])
answer = output_text.split("assistant<|end_of_role|>")[1]
print("Answer: " + answer)

# Generate requirement check
req_generation_prompt = "<|start_of_role|>check_requirement<|end_of_role|>"
req_chat = question_chat + [
    {
        "role": "assistant",
        "content": answer
    },
]

req_text = tokenizer.apply_chat_template(req_chat,tokenize=False) 
req_text = req_text[len_sys:]
# tokenize and generate
inputs, alora_offsets = tokenize_alora(tokenizer,req_text, req_generation_prompt)

if REUSE_CACHE: #reuse KV cache from earlier answer generation 
    output = model_req_check.generate(inputs["input_ids"].to(device), attention_mask=inputs["attention_mask"].to(device), max_new_tokens=1,alora_offsets=alora_offsets,past_key_values=answer_cache)
else: #simplest call
    output = model_req_check.generate(inputs["input_ids"].to(device), attention_mask=inputs["attention_mask"].to(device), max_new_tokens=1,alora_offsets=alora_offsets)
output_text = tokenizer.decode(output[0])

# Extract score
req_score = output_text[-1]
print("Requirement Satisfied: " + req_score)