mirror of
https://github.com/SakanaAI/doc-to-lora.git
synced 2026-07-23 17:01:04 +02:00
137 lines
4.5 KiB
Python
137 lines
4.5 KiB
Python
import json
|
|
|
|
import torch
|
|
from modeling_icae_multi_span import (
|
|
ICAE,
|
|
DataArguments,
|
|
ModelArguments,
|
|
TrainingArguments,
|
|
)
|
|
from peft import LoraConfig
|
|
from safetensors.torch import load_file
|
|
from tqdm import tqdm
|
|
from transformers import HfArgumentParser
|
|
|
|
# Set the computation device
|
|
device = "cuda"
|
|
|
|
# Parse model, data, and training arguments
|
|
parser = HfArgumentParser((ModelArguments, DataArguments, TrainingArguments))
|
|
model_args, data_args, training_args = parser.parse_args_into_dataclasses()
|
|
|
|
# Define Lora configuration
|
|
lora_config = LoraConfig(
|
|
r=512,
|
|
lora_alpha=32,
|
|
lora_dropout=model_args.lora_dropout,
|
|
bias="none",
|
|
task_type="CAUSAL_LM",
|
|
)
|
|
|
|
# Initialize model and send it to CUDA device
|
|
model = ICAE(model_args, training_args, lora_config)
|
|
|
|
# Load the fine-tuned checkpoint
|
|
print(f"Loading trained checkpoint from {training_args.output_dir}")
|
|
state_dict = load_file(training_args.output_dir)
|
|
model.load_state_dict(
|
|
state_dict, strict=False
|
|
) # only load lora and memory token embeddings
|
|
|
|
model = model.to(device)
|
|
|
|
# Read the data file
|
|
file_path = "./dev_v2.jsonl"
|
|
lines = None
|
|
with open(file_path) as f:
|
|
lines = f.readlines()
|
|
|
|
# Prepare the model for evaluation
|
|
max_out_length = 512
|
|
model.eval()
|
|
|
|
with torch.no_grad():
|
|
with open("ft_inference.out", "w") as f:
|
|
for line in tqdm(lines):
|
|
# Tokenize input text
|
|
data = json.loads(line)
|
|
tokenized_input = model.tokenizer(
|
|
data["input"],
|
|
truncation=True,
|
|
max_length=5120,
|
|
padding=False,
|
|
return_attention_mask=False,
|
|
)
|
|
tokenized_prompt = model.tokenizer(
|
|
data["prompt"],
|
|
truncation=False,
|
|
padding=False,
|
|
return_attention_mask=False,
|
|
add_special_tokens=False,
|
|
)
|
|
# Generate compressed outputs
|
|
input_ids = torch.LongTensor([tokenized_input["input_ids"]]).to(device)
|
|
memory_slots = model._compress(input_ids)
|
|
|
|
# decoder input has 3 parts: prefix, memory slots and suffix
|
|
# the following code is for Mistral tokenizer for example: 733, 16289, 28793 are for the Mistral instruction tempmlate
|
|
prompt_left_ids = torch.LongTensor([[1, 733, 16289, 28793]]).to(device)
|
|
prompt_right_ids = (
|
|
[model.ft_token_id]
|
|
+ tokenized_prompt["input_ids"]
|
|
+ [733, 28748, 16289, 28793]
|
|
)
|
|
prompt_right_ids = torch.LongTensor([prompt_right_ids]).to(device)
|
|
|
|
prompt_left_embs = model.tokens_to_embeddings(prompt_left_ids)
|
|
prompt_right_embs = model.tokens_to_embeddings(prompt_right_ids)
|
|
memory_slots = memory_slots.to(prompt_right_embs)
|
|
|
|
# Concatenate and clone input embeddings
|
|
decoder_input_embeddings = torch.cat(
|
|
(prompt_left_embs, memory_slots.unsqueeze(0), prompt_right_embs), dim=1
|
|
)
|
|
output = decoder_input_embeddings.clone()
|
|
|
|
generate_text = []
|
|
past_key_values = None
|
|
|
|
# Generate text output
|
|
for i in range(max_out_length):
|
|
with (
|
|
model.icae.disable_adapter()
|
|
): # no independent decoder; use self.icae
|
|
out = model.icae(
|
|
inputs_embeds=output,
|
|
past_key_values=past_key_values,
|
|
use_cache=True,
|
|
)
|
|
# out = decoder(inputs_embeds=output, past_key_values=past_key_values, use_cache=True)
|
|
logit = out.logits[:, -1, : model.vocab_size - 1]
|
|
past_key_values = out.past_key_values
|
|
|
|
next_token_id = torch.argmax(logit, dim=-1)
|
|
# print(next_token_id)
|
|
|
|
if next_token_id.item() == 2: # eos
|
|
break
|
|
|
|
output = (
|
|
model.icae.get_base_model()
|
|
.model.embed_tokens(next_token_id)
|
|
.unsqueeze(1)
|
|
.to(device)
|
|
)
|
|
generate_text.append(next_token_id.item())
|
|
|
|
generated_text = model.tokenizer.decode(generate_text)
|
|
|
|
# Structure output data
|
|
output_ = {
|
|
"input": data["input"],
|
|
"prompt": data["prompt"],
|
|
"output": generated_text,
|
|
"answer": data["answer"],
|
|
}
|
|
|
|
f.write(json.dumps(output_) + "\n")
|