| import re |
| from typing import List |
|
|
| import click |
| import torch |
|
|
|
|
| def prepare_results(input_data: List[str], predictions: List[dict]) -> List[List[str]]: |
| return [ |
| phla.split(",") + [str(pred["Presentation score"]), pred["Label"]] |
| for phla, pred in zip(input_data, predictions) |
| ] |
|
|
|
|
| @click.command( |
| help="This is a model for prediction the presentation of a peptide on MHC class 1. " |
| "The model runs on inputs filtered by mhcflurry (and includes the filtering). " |
| "The model was created as part of Immunology platform of Ardigen. " |
| "For more information, see https://huggingface.co/ardigen/ardisplay-i" |
| ) |
| @click.option( |
| "--input-file", |
| "-i", |
| type=click.Path(exists=True), |
| required=True, |
| help="Path to input CSV file. The file should contain one peptide-HLA pair per line," |
| " with the HLA and peptide separated by a comma (e.g. A02:01,CKTSPLSNWHT)'. " |
| "Note that file should not contain any headers.", |
| ) |
| @click.option( |
| "--output-file", |
| "-o", |
| default="ardisplay_predictions.csv", |
| type=click.Path(), |
| help="Path to output CSV file. Default is ardisplay_predictions.csv", |
| ) |
| @click.option( |
| "--batch-size", |
| "-b", |
| type=int, |
| default=100, |
| help="Batch size for inference. Default is 100.", |
| ) |
| @click.option( |
| "--device", |
| "-d", |
| type=str, |
| default="cpu", |
| help="Device to use for inference, either 'cpu' or a CUDA device (e.g. 'cuda:0'). Default is 'cpu'.", |
| ) |
| def ardisplay_cli(input_file, output_file, batch_size, device): |
| from datasets import load_dataset |
| from transformers import pipeline |
|
|
| if not re.match("^cpu$|^cuda:\d+$", device): |
| raise click.BadParameter( |
| "Device should be 'cpu' or a CUDA device (e.g. 'cuda:0')" |
| ) |
|
|
| device = torch.device(device) |
| pipe = pipeline( |
| model="ardigen/ardisplay-i", |
| trust_remote_code=True, |
| batch_size=batch_size, |
| device=device, |
| ) |
| dataset = load_dataset("text", data_files={"test": input_file}, split="test") |
| predictions = pipe(dataset["text"]) |
|
|
| results = prepare_results(input_data=dataset["text"], predictions=predictions) |
| with open(output_file, "w") as f: |
| f.write("HLA,Peptide,Presentation score,Label\n") |
| for result in results: |
| f.write(",".join(result) + "\n") |
| click.echo(f"Presentation scores calculated, outputting to {output_file}") |
|
|
|
|
| if __name__ == "__main__": |
| ardisplay_cli() |
|
|
