PTMPicker: Facilitating Efficient Pretrained Model Selection for Application Developers

TL;DR

PTMPicker is a tool that improves pretrained model selection by representing models and user requirements in a structured format, enabling more accurate and efficient identification of suitable models considering various constraints.

Contribution

It introduces a structured attribute-based representation for models and search requests, and employs similarity and prompt-based evaluation to enhance model selection accuracy.

Findings

Achieved 85% success rate in top-10 model retrieval.

Scraped and processed over 543,000 models from Hugging Face.

Generated 15,207 synthetic search requests for evaluation.

Abstract

The rapid emergence of pretrained models (PTMs) has attracted significant attention from both Deep Learning (DL) researchers and downstream application developers. However, selecting appropriate PTMs remains challenging because existing methods typically rely on keyword-based searches in which the keywords are often derived directly from function descriptions. This often fails to fully capture user intent and makes it difficult to identify suitable models when developers also consider factors such as bias mitigation, hardware requirements, or license compliance. To address the limitations of keyword-based model search, we propose PTMPicker to accurately identify suitable PTMs. We first define a structured template composed of common and essential attributes for PTMs and then PTMPicker represents both candidate models and user-intended features (i.e., model search requests) in this unified format. To determine whether candidate models satisfy user requirements, it computes embedding similarities for function-related attributes and uses well-crafted prompts to evaluate special constraints such as license compliance and hardware requirements. We scraped a total of 543,949 pretrained models from Hugging Face to prepare valid candidates for selection. PTMPicker then represented them in the predefined structured format by extracting their associated descriptions. Guided by the extracted metadata, we synthesized a total of 15,207 model search requests with carefully designed prompts, as no such search requests are readily available. Experiments on the curated PTM dataset and the synthesized model search requests show that PTMPicker can help users effectively identify models,with 85% of the sampled requests successfully locating appropriate PTMs within the top-10 ranked candidates.