🧱 Integrate SetFit library (#387)

tomaarsen · web-flow · commit d724ce3c8497 · 2023-12-05T23:33:24.000+01:00
Hello! ## Pull Request overview * Integrate with the [SetFit](https://github.com/huggingface/setfit) library for Text Classification. ## Details [SetFit](https://github.com/huggingface/setfit) is a library for text classification with ~1200 models on the Hub at the time of writing. A v1.0.0 release is upcoming, and it's a good time to add this widget support, etc. It can be used like so: ```python from setfit import SetFitModel model = SetFitModel.from_pretrained("tomaarsen/span-marker-bert-base-fewnerd-fine-super") ``` ``` model.predict(["That was an awful movie"]) # => ["negative"] ``` I've previously integrated a library by only editing hub-docs and api-inference-community, but I see that there's been some refactors since. I hope that with these changes, I've edited the correct places. I also noticed this file: https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/src/library-to-tasks.ts#L36, but it seems that it's automatically updated. So, I didn't touch that one. Let me know if there's any more changes needed! Related PRs: * huggingface/api-inference-community#359 * huggingface/hub-docs#1150 - Tom Aarsen
diff --git a/packages/tasks/src/library-to-tasks.ts b/packages/tasks/src/library-to-tasks.ts
@@ -31,6 +31,7 @@ export const LIBRARY_TASK_MAPPING_EXCLUDING_TRANSFORMERS: Partial<Record<ModelLi
 	peft: ["text-generation"],
 	"pyannote-audio": ["automatic-speech-recognition"],
 	"sentence-transformers": ["feature-extraction", "sentence-similarity"],
+	setfit: ["text-classification"],
 	sklearn: ["tabular-classification", "tabular-regression", "text-classification"],
 	spacy: ["token-classification", "text-classification", "sentence-similarity"],
 	"span-marker": ["token-classification"],
diff --git a/packages/tasks/src/library-ui-elements.ts b/packages/tasks/src/library-ui-elements.ts
@@ -358,6 +358,12 @@ const sentenceTransformers = (model: ModelData) => [
 model = SentenceTransformer("${model.id}")`,
 ];
 
+const setfit = (model: ModelData) => [
+	`from setfit import SetFitModel
+
+model = SetFitModel.from_pretrained("${model.id}")`,
+];
+
 const spacy = (model: ModelData) => [
 	`!pip install https://huggingface.co/${model.id}/resolve/main/${nameWithoutNamespace(model.id)}-any-py3-none-any.whl
 
@@ -661,6 +667,13 @@ export const MODEL_LIBRARIES_UI_ELEMENTS: Partial<Record<ModelLibraryKey, Librar
 		docsUrl: "https://huggingface.co/docs/hub/sentence-transformers",
 		snippets: sentenceTransformers,
 	},
+	setfit: {
+		btnLabel: "setfit",
+		repoName: "setfit",
+		repoUrl: "https://github.com/huggingface/setfit",
+		docsUrl: "https://huggingface.co/docs/hub/setfit",
+		snippets: setfit,
+	},
 	sklearn: {
 		btnLabel: "Scikit-learn",
 		repoName: "Scikit-learn",
diff --git a/packages/tasks/src/model-libraries.ts b/packages/tasks/src/model-libraries.ts
@@ -23,6 +23,7 @@ export enum ModelLibrary {
 	"pyannote-audio" = "pyannote.audio",
 	"sample-factory" = "Sample Factory",
 	"sentence-transformers" = "Sentence Transformers",
+	"setfit" = "SetFit",
 	"sklearn" = "Scikit-learn",
 	"spacy" = "spaCy",
 	"span-marker" = "SpanMarker",
diff --git a/packages/tasks/src/tasks/index.ts b/packages/tasks/src/tasks/index.ts
@@ -69,7 +69,7 @@ export const TASKS_MODEL_LIBRARIES: Record<PipelineType, ModelLibraryKey[]> = {
 	"tabular-classification": ["sklearn"],
 	"tabular-regression": ["sklearn"],
 	"tabular-to-text": ["transformers"],
-	"text-classification": ["adapter-transformers", "spacy", "transformers", "transformers.js"],
+	"text-classification": ["adapter-transformers", "setfit", "spacy", "transformers", "transformers.js"],
 	"text-generation": ["transformers", "transformers.js"],
 	"text-retrieval": [],
 	"text-to-image": ["diffusers"],
