Spaces:
Running
Running
File size: 4,117 Bytes
08ccc8e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 |
import argparse
import os
import sys
import warnings
import numpy as np
import pandas as pd
import torch
from torch.utils.data import DataLoader
from transformers import AutoTokenizer, T5EncoderModel
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
from generation_utils import ReactionT5Dataset
from train import preprocess_df, preprocess_USPTO
from utils import filter_out, seed_everything
warnings.filterwarnings("ignore")
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument(
"--input_data",
type=str,
required=True,
help="Path to the input data.",
)
parser.add_argument(
"--test_data",
type=str,
required=False,
help="Path to the test data. If provided, the duplicates will be removed from the input data.",
)
parser.add_argument(
"--input_max_length",
type=int,
default=400,
help="Maximum token length of input.",
)
parser.add_argument(
"--model_name_or_path",
type=str,
default="sagawa/ReactionT5v2-forward",
help="Name or path of the finetuned model for prediction. Can be a local model or one from Hugging Face.",
)
parser.add_argument(
"--batch_size", type=int, default=5, help="Batch size for prediction."
)
parser.add_argument(
"--output_dir",
type=str,
default="./",
help="Directory where predictions are saved.",
)
parser.add_argument(
"--debug", action="store_true", default=False, help="Use debug mode."
)
parser.add_argument(
"--seed", type=int, default=42, help="Seed for reproducibility."
)
return parser.parse_args()
def create_embedding(dataloader, model, device):
outputs_mean = []
model.eval()
model.to(device)
for inputs in dataloader:
inputs = {k: v.to(CFG.device) for k, v in inputs.items()}
with torch.no_grad():
output = model(**inputs)
last_hidden_states = output[0]
input_mask_expanded = (
inputs["attention_mask"]
.unsqueeze(-1)
.expand(last_hidden_states.size())
.float()
)
sum_embeddings = torch.sum(last_hidden_states * input_mask_expanded, 1)
sum_mask = input_mask_expanded.sum(1)
sum_mask = torch.clamp(sum_mask, min=1e-6)
mean_embeddings = sum_embeddings / sum_mask
outputs_mean.append(mean_embeddings.detach().cpu().numpy())
return outputs_mean
if __name__ == "__main__":
CFG = parse_args()
CFG.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
if not os.path.exists(CFG.output_dir):
os.makedirs(CFG.output_dir)
seed_everything(seed=CFG.seed)
CFG.tokenizer = AutoTokenizer.from_pretrained(
os.path.abspath(CFG.model_name_or_path)
if os.path.exists(CFG.model_name_or_path)
else CFG.model_name_or_path,
return_tensors="pt",
)
model = T5EncoderModel.from_pretrained(CFG.model_name_or_path).to(CFG.device)
model.eval()
input_data = filter_out(pd.read_csv(CFG.input_data), ["REACTANT", "PRODUCT"])
input_data = preprocess_df(input_data, drop_duplicates=False)
if CFG.test_data:
input_data_copy = preprocess_USPTO(input_data.copy())
test_data = filter_out(pd.read_csv(CFG.test_data), ["REACTANT", "PRODUCT"])
USPTO_test = preprocess_USPTO(test_data)
input_data = input_data[
~input_data_copy["pair"].isin(USPTO_test["pair"])
].reset_index(drop=True)
input_data.to_csv(os.path.join(CFG.output_dir, "input_data.csv"), index=False)
dataset = ReactionT5Dataset(CFG, input_data)
dataloader = DataLoader(
dataset,
batch_size=CFG.batch_size,
shuffle=False,
num_workers=4,
pin_memory=True,
drop_last=False,
)
outputs = create_embedding(dataloader, model, CFG.device)
outputs = np.concatenate(outputs, axis=0)
np.save(os.path.join(CFG.output_dir, "embedding_mean.npy"), outputs)
|