A scikit-learn API version of a SetFit classifier. Model originally developed by Moshe Wasserblat.
from setfit import SetFitClassifier
docs = ["yay", "boo", "yes", "no", "yeah"]
labels = [1, 0, 1, 0, 1]
# takes a sentence-transformers model
clf = SetFitClassifier("paraphrase-MiniLM-L3-v2")
# fine-tunes embeddings + trains logistic regression head
clf.fit(docs, labels)
clf.predict(["affirmitive", "negative"])
array([1, 0])pip install git+https://github.com/pmbaumgartner/setfit
Reference Notebook:
