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.
翻译:预训练模型(PTMs)的快速涌现已引起深度学习(DL)研究者和下游应用开发者的广泛关注。然而,选择合适的预训练模型仍然具有挑战性,因为现有方法通常依赖于基于关键词的搜索,而这些关键词往往直接来自功能描述。这常常无法充分捕捉用户意图,并且当开发者还需考虑偏差缓解、硬件要求或许可证合规性等因素时,难以识别合适的模型。为克服基于关键词的模型搜索的局限性,我们提出PTMPicker以精准识别合适的预训练模型。我们首先定义了一个由预训练模型的通用及核心属性组成的结构化模板,随后PTMPicker将候选模型和用户期望特征(即模型搜索请求)均以此统一格式表示。为判断候选模型是否满足用户需求,该方法计算功能相关属性的嵌入相似度,并利用精心设计的提示词来评估许可证合规性和硬件要求等特殊约束。我们从Hugging Face平台共爬取了543,949个预训练模型作为有效候选集。PTMPicker通过提取相关描述,将这些模型以预定义的结构化格式进行表示。在提取的元数据指导下,我们合成了总计15,207个模型搜索请求(采用精心设计的提示词),因为目前尚无此类现成的搜索请求可用。在精心整理的预训练模型数据集和合成的模型搜索请求上进行的实验表明,PTMPicker能有效帮助用户识别模型,在抽样请求中,85%的请求能在排名前10的候选模型中成功定位到合适的预训练模型。