Automated categorization of pre-trained models for software engineering: A case study with a Hugging Face dataset

TL;DR

This paper proposes an automated method to categorize pre-trained models for software engineering tasks on the Hugging Face platform, aiding users in selecting appropriate models through a similarity-based classification approach.

Contribution

It introduces a semi-automated approach to classify PTMs for SE tasks by leveraging Hugging Face data and literature, filling a gap in existing model categorization.

Findings

Model cards contain sufficient information for classification.

The approach effectively maps PTMs to specific SE tasks.

The method improves model selection for SE applications.

Abstract

Software engineering (SE) activities have been revolutionized by the advent of pre-trained models (PTMs), defined as large machine learning (ML) models that can be fine-tuned to perform specific SE tasks. However, users with limited expertise may need help to select the appropriate model for their current task. To tackle the issue, the Hugging Face (HF) platform simplifies the use of PTMs by collecting, storing, and curating several models. Nevertheless, the platform currently lacks a comprehensive categorization of PTMs designed specifically for SE, i.e., the existing tags are more suited to generic ML categories. This paper introduces an approach to address this gap by enabling the automatic classification of PTMs for SE tasks. First, we utilize a public dump of HF to extract PTMs information, including model documentation and associated tags. Then, we employ a semi-automated method to identify SE tasks and their corresponding PTMs from existing literature. The approach involves creating an initial mapping between HF tags and specific SE tasks, using a similarity-based strategy to identify PTMs with relevant tags. The evaluation shows that model cards are informative enough to classify PTMs considering the pipeline tag. Moreover, we provide a mapping between SE tasks and stored PTMs by relying on model names.