Repository: refuel-ai/autolabel

Issue:
[Bug]: Attempting to use provider Google with "few_shot_selection": "semantic_similarity" raises 'VertexAIEmbeddings' object has no attribute 'model'
When attempting to use provider "google" with model "text-bison@001" and "few_shot_selection": "semantic_similarity" raises 'VertexAIEmbeddings' object has no attribute 'model'.

Full stack trace:
```
File ~/anaconda3/envs/text_gen/lib/python3.10/site-packages/autolabel/labeler.py:325, in LabelingAgent.plan(self, dataset, max_items, start_index)
    316 if (
    317     self.config.explanation_column()
    318     and len(seed_examples) > 0
    319     and self.config.explanation_column() not in list(seed_examples[0].keys())
    320 ):
    321     raise ValueError(
    322         f"Explanation column {self.config.explanation_column()} not found in dataset.\nMake sure that explanations were generated using labeler.generate_explanations(seed_file)."
    323     )
--> 325 self.example_selector = ExampleSelectorFactory.initialize_selector(
    326     self.config,
    327     seed_examples,
    328     dataset.df.keys().tolist(),
    329     cache=self.generation_cache is not None,
    330 )
    332 input_limit = min(len(dataset.inputs), 100)
    334 for input_i in track(
    335     dataset.inputs[:input_limit],
    336     description="Generating Prompts...",
    337     console=self.console,
    338 ):
    339     # TODO: Check if this needs to use the example selector

File ~/anaconda3/envs/text_gen/lib/python3.10/site-packages/autolabel/few_shot/__init__.py:115, in ExampleSelectorFactory.initialize_selector(config, examples, columns, cache)
    109 if algorithm not in [
    110     FewShotAlgorithm.FIXED,
    111     FewShotAlgorithm.LABEL_DIVERSITY_RANDOM,
    112 ]:
    113     params["cache"] = cache
--> 115 return example_cls.from_examples(**params)

File ~/anaconda3/envs/text_gen/lib/python3.10/site-packages/langchain/prompts/example_selector/semantic_similarity.py:96, in SemanticSimilarityExampleSelector.from_examples(cls, examples, embeddings, vectorstore_cls, k, input_keys, **vectorstore_cls_kwargs)
     94 else:
     95     string_examples = [" ".join(sorted_values(eg)) for eg in examples]
---> 96 vectorstore = vectorstore_cls.from_texts(
     97     string_examples, embeddings, metadatas=examples, **vectorstore_cls_kwargs
     98 )
     99 return cls(vectorstore=vectorstore, k=k, input_keys=input_keys)

File ~/anaconda3/envs/text_gen/lib/python3.10/site-packages/autolabel/few_shot/vector_store.py:453, in VectorStoreWrapper.from_texts(cls, texts, embedding, metadatas, cache, **kwargs)
    436 """Create a vectorstore from raw text.
    437 The data will be ephemeral in-memory.
    438 Args:
   (...)
    444     vector_store: Vectorstore with seedset embeddings
    445 """
    446 vector_store = cls(
    447     embedding_function=embedding,
    448     corpus_embeddings=None,
   (...)
    451     **kwargs,
    452 )
--> 453 vector_store.add_texts(texts=texts, metadatas=metadatas)
    454 return vector_store

File ~/anaconda3/envs/text_gen/lib/python3.10/site-packages/autolabel/few_shot/vector_store.py:244, in VectorStoreWrapper.add_texts(self, texts, metadatas)
    236 """Run texts through the embeddings and add to the vectorstore. Currently, the vectorstore is reinitialized each time, because we do not require a persistent vector store for example selection.
    237 Args:
    238     texts (Iterable[str]): Texts to add to the vectorstore.
   (...)
    241     List[str]: List of IDs of the added texts.
    242 """
    243 if self._embedding_function is not None:
--> 244     embeddings = self._get_embeddings(texts)
    246 self._corpus_embeddings = torch.tensor(embeddings)
    247 self._texts = texts

File ~/anaconda3/envs/text_gen/lib/python3.10/site-packages/autolabel/few_shot/vector_store.py:181, in VectorStoreWrapper._get_embeddings(self, texts)
    176 for idx, text in enumerate(texts):
    177     query = sql_text(
    178         f"SELECT embedding FROM {EMBEDDINGS_TABLE} WHERE embedding_function = :x AND text = :y",
    179     )
    180     params = {
--> 181         "x": self._embedding_function.model
    182         if self._embedding_function.__class__.__name__
    183         != "HuggingFaceEmbeddings"
    184         else self._embedding_function.model_name,
    185         "y": text,
    186     }
    187     result = conn.execute(query, params).fetchone()
    188     if result:

AttributeError: 'VertexAIEmbeddings' object has no attribute 'model'
```


Discussion:

Resolving patch:
diff --git a/src/autolabel/few_shot/vector_store.py b/src/autolabel/few_shot/vector_store.py
--- a/src/autolabel/few_shot/vector_store.py
+++ b/src/autolabel/few_shot/vector_store.py
@@ -180,7 +180,7 @@ def _get_embeddings(self, texts: Iterable[str]) -> List[List[float]]:
                     params = {
                         "x": self._embedding_function.model
                         if self._embedding_function.__class__.__name__
-                        != "HuggingFaceEmbeddings"
+                        not in ["HuggingFaceEmbeddings", "VertexAIEmbeddings"]
                         else self._embedding_function.model_name,
                         "y": text,
                     }
@@ -220,7 +220,7 @@ def _add_embeddings_to_cache(
                     params = {
                         "x": self._embedding_function.model
                         if self._embedding_function.__class__.__name__
-                        != "HuggingFaceEmbeddings"
+                        not in ["HuggingFaceEmbeddings", "VertexAIEmbeddings"]
                         else self._embedding_function.model_name,
                         "y": text,
                         "z": pickle.dumps(embedding),


Test patch: