Spaces:
Running
Running
import argparse | |
import logging | |
import os | |
import sys | |
import warnings | |
import pandas as pd | |
import torch | |
from datasets.utils.logging import disable_progress_bar | |
from torch.utils.data import DataLoader | |
from transformers import AutoTokenizer | |
# Suppress warnings and logging | |
warnings.filterwarnings("ignore") | |
logging.disable(logging.WARNING) | |
disable_progress_bar() | |
os.environ["TOKENIZERS_PARALLELISM"] = "false" | |
# Append the utils module path | |
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) | |
from generation_utils import ReactionT5Dataset | |
from models import ReactionT5Yield2 | |
from prediction import inference_fn | |
from train import preprocess_df | |
from utils import seed_everything | |
def parse_args(): | |
""" | |
Parse command line arguments. | |
""" | |
parser = argparse.ArgumentParser( | |
description="Prediction script for ReactionT5Yield model." | |
) | |
parser.add_argument( | |
"--input_data", | |
type=str, | |
required=True, | |
help="Data as a CSV file that contains an 'input' column. The format of the contents of the column are like 'REACTANT:{reactants of the reaction}PRODUCT:{products of the reaction}'. If there are multiple reactants, concatenate them with '.'.", | |
) | |
parser.add_argument( | |
"--model_name_or_path", | |
type=str, | |
default="sagawa/ReactionT5v2-yield", | |
help="Name or path of the finetuned model for prediction. Can be a local model or one from Hugging Face.", | |
) | |
parser.add_argument("--debug", action="store_true", help="Use debug mode.") | |
parser.add_argument( | |
"--input_max_length", | |
type=int, | |
default=400, | |
help="Maximum token length of input.", | |
) | |
parser.add_argument( | |
"--batch_size", type=int, default=5, required=False, help="Batch size." | |
) | |
parser.add_argument( | |
"--num_workers", type=int, default=4, help="Number of data loading workers." | |
) | |
parser.add_argument( | |
"--fc_dropout", | |
type=float, | |
default=0.0, | |
help="Dropout rate after fully connected layers.", | |
) | |
parser.add_argument( | |
"--output_dir", | |
type=str, | |
default="./", | |
help="Directory where predictions are saved.", | |
) | |
parser.add_argument( | |
"--seed", type=int, default=42, help="Random seed for reproducibility." | |
) | |
return parser.parse_args() | |
if __name__ == "__main__": | |
CFG = parse_args() | |
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
CFG.device = device | |
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 = ReactionT5Yield2.from_pretrained(CFG.model_name_or_path) | |
test_ds = pd.read_csv(CFG.input_data) | |
test_ds = preprocess_df(test_ds, CFG, drop_duplicates=False) | |
test_dataset = ReactionT5Dataset(CFG, test_ds) | |
test_loader = DataLoader( | |
test_dataset, | |
batch_size=CFG.batch_size, | |
shuffle=False, | |
num_workers=CFG.num_workers, | |
pin_memory=True, | |
drop_last=False, | |
) | |
prediction = inference_fn(test_loader, model, CFG) | |
test_ds["prediction"] = prediction | |
test_ds["prediction"] = test_ds["prediction"].clip(0, 100) | |
test_ds.to_csv( | |
os.path.join(CFG.output_dir, "yield_prediction_output.csv"), index=False | |
) | |