Added better functionality for label selection

src/autolabel/configs/config.py

@@ -237,7 +237,7 @@ def label_selection(self) -> bool:
         classification tasks with a large number of possible classes."""
         return self._prompt_config.get(self.LABEL_SELECTION_KEY, False)
 
-    def label_selection_count(self) -> int:
+    def max_selected_labels(self) -> int:
         """Returns the number of labels to select in LabelSelector"""
         k = self._prompt_config.get(self.LABEL_SELECTION_COUNT_KEY, 10)
         if k < 1:
@@ -248,7 +248,7 @@ def label_selection_threshold(self) -> float:
         """Returns the threshold for label selection in LabelSelector
         If the similarity score ratio with the top Score is above this threshold,
         the label is selected."""
-        return self._prompt_config.get(self.LABEL_SELECTION_THRESHOLD, 0.95)
+        return self._prompt_config.get(self.LABEL_SELECTION_THRESHOLD, 0.0)
 
     def attributes(self) -> List[Dict]:
         """Returns a list of attributes to extract from the text."""

src/autolabel/few_shot/label_selector.py

@@ -4,6 +4,7 @@
 from collections.abc import Callable
 from typing import Dict, List, Optional, Tuple, Union
 
+import torch
 from sqlalchemy.sql import text as sql_text
 
 from autolabel.configs import AutolabelConfig
@@ -36,7 +37,7 @@ def __init__(
         self.config = config
         self.labels = self.config.labels_list()
         self.label_descriptions = self.config.label_descriptions()
-        self.k = min(self.config.label_selection_count(), len(self.labels))
+        self.k = min(self.config.max_selected_labels(), len(self.labels))
         self.threshold = self.config.label_selection_threshold()
         self.cache = cache
         self.vectorStore = VectorStoreWrapper(
@@ -46,23 +47,22 @@ def __init__(
         # Get the embeddings of the labels
         if self.label_descriptions is not None:
             (labels, descriptions) = zip(*self.label_descriptions.items())
-            embeddings = self.vectorStore._get_embeddings(descriptions)
-            for i, label in enumerate(labels):
-                self.labels_embeddings[label] = embeddings[i]
+            self.labels = list(labels)
+            self.labels_embeddings = torch.Tensor(
+                self.vectorStore._get_embeddings(descriptions)
+            )
         else:
-            embeddings = self.vectorStore._get_embeddings(self.labels)
-            for i, label in enumerate(labels):
-                self.labels_embeddings[label] = embeddings[i]
+            self.labels_embeddings = torch.Tensor(
+                self.vectorStore._get_embeddings(self.labels)
+            )
+        print(type(self.labels_embeddings))
 
     def select_labels(self, input: str) -> List[str]:
         """Select which labels to use based on the similarity to input"""
-        input_embedding = self.vectorStore._get_embeddings([input])
-
-        scores = []
-        for label, embedding in self.labels_embeddings.items():
-            similarity = cos_sim(embedding, input_embedding)
-            # insert into scores, while maintaining sorted order
-            bisect.insort(scores, (similarity, label))
+        input_embedding = torch.Tensor(self.vectorStore._get_embeddings([input]))
+        scores = cos_sim(input_embedding, self.labels_embeddings).view(-1)
+        scores = list(zip(scores, self.labels))
+        scores.sort(key=lambda x: x[0])
 
         # remove labels with similarity score less than self.threshold*topScore
         return [

src/autolabel/labeler.py

@@ -226,6 +226,7 @@ def run(
                 self.config.label_selection()
                 and self.config.few_shot_algorithm() != "fixed"
             ):
+                # TODO: Add support for other few shot algorithms specially semantic similarity
                 raise ValueError(
                     "Error: Only 'fixed' few shot example selector is supported for label selection."
                 )
@@ -271,7 +272,8 @@ def run(
             ):
                 # get every column except the one we want to label
                 toEmbed = chunk.copy()
-                del toEmbed[self.config.label_column()]
+                if self.config.label_column():
+                    del toEmbed[self.config.label_column()]
 
                 # convert this to a string
                 toEmbed = json.dumps(toEmbed)