vertex zero shot

Co-authored-by: Iryna Kondrashchenko <[email protected]>

Author: OKUA1

skllm/llm/vertex/mixin.py

@@ -0,0 +1,73 @@
+from typing import Optional, Union, List, Any, Dict, Mapping
+from skllm.config import SKLLMConfig as _Config
+from skllm.llm.base import (
+    BaseClassifierMixin,
+    BaseEmbeddingMixin,
+    BaseTextCompletionMixin,
+    BaseTunableMixin,
+)
+from skllm.llm.vertex.completion import get_completion_chat_mode, get_completion
+from skllm.utils import extract_json_key
+import numpy as np
+from tqdm import tqdm
+import json
+
+
+class VertexMixin:
+    pass
+
+
+class VertexTextCompletionMixin(BaseTextCompletionMixin):
+    def _get_chat_completion(
+        self,
+        model: str,
+        messages: Union[str, List[Dict[str, str]]],
+        system_message: Optional[str],
+        examples: Optional[List] = None,
+    ) -> str:
+        if examples is not None:
+            raise NotImplementedError(
+                "Examples API is not yet supported for Vertex AI."
+            )
+        if not isinstance(messages, str):
+            raise ValueError("Only messages as strings are supported.")
+        if model.startswith("chat-"):
+            completion = get_completion_chat_mode(model, system_message, messages)
+        else:
+            completion = get_completion(model, messages)
+        return str(completion)
+
+
+class VertexClassifierMixin(BaseClassifierMixin, VertexTextCompletionMixin):
+    def _extract_out_label(self, completion: str, **kwargs) -> Any:
+        """Extracts the label from a completion.
+
+        Parameters
+        ----------
+        label : Mapping[str, Any]
+            The label to extract.
+
+        Returns
+        -------
+        label : str
+        """
+        print(completion)
+        try:
+            label = extract_json_key(str(completion), "label")
+        except Exception as e:
+            print(completion)
+            print(f"Could not extract the label from the completion: {str(e)}")
+            label = ""
+        print(label)
+        return label
+
+
+class VertexEmbeddingMixin(BaseEmbeddingMixin):
+    def _get_embeddings(self, text: np.ndarray) -> List[List[float]]:
+        raise NotImplementedError("Embeddings are not yet supported for Vertex AI.")
+
+
+class VertexTunableMixin(BaseTunableMixin):
+    # TODO
+    def _tune(self, X: Any, y: Any):
+        raise NotImplementedError("Tuning is not yet supported for Vertex AI.")

skllm/models/_base/classifier.py

@@ -31,6 +31,7 @@
 from skllm.memory.base import IndexConstructor
 from skllm.memory._sklearn_nn import SklearnMemoryIndex
 from skllm.models._base.vectorizer import BaseVectorizer as _BaseVectorizer
+import ast
 
 _TRAINING_SAMPLE_PROMPT_TEMPLATE = """
 Sample input:

skllm/models/vertex/classification/zero_shot.py

@@ -0,0 +1,45 @@
+from skllm.llm.vertex.mixin import VertexClassifierMixin as _VertexClassifierMixin
+from skllm.models._base.classifier import (
+    BaseZeroShotClassifier as _BaseZeroShotClassifier,
+    SingleLabelMixin as _SingleLabelMixin,
+    MultiLabelMixin as _MultiLabelMixin,
+)
+from typing import Optional
+
+
+class ZeroShotVertexClassifier(
+    _BaseZeroShotClassifier, _SingleLabelMixin, _VertexClassifierMixin
+):
+    def __init__(
+        self,
+        model: str = "text-bison@001",
+        default_label: Optional[str] = "Random",
+        prompt_template: Optional[str] = None,
+        **kwargs,
+    ):
+        super().__init__(
+            model=model,
+            default_label=default_label,
+            prompt_template=prompt_template,
+            **kwargs,
+        )
+
+
+class MultiLabelZeroShotVertexClassifier(
+    _BaseZeroShotClassifier, _MultiLabelMixin, _VertexClassifierMixin
+):
+    def __init__(
+        self,
+        model: str = "text-bison@001",
+        default_label: Optional[str] = "Random",
+        prompt_template: Optional[str] = None,
+        max_labels: Optional[int] = 5,
+        **kwargs,
+    ):
+        super().__init__(
+            model=model,
+            default_label=default_label,
+            prompt_template=prompt_template,
+            max_labels=max_labels,
+            **kwargs,
+        )