[Bug]: Gemini failed in Vertex AI workbench

[LeoMai2024]

Describe the bug Hi AutoLabel team, Amazing project! Currently, I am trying to use autolabel in Vertex AI workbench, I successfully installed autolabel library and I have Gemini pro access. However, the notebook will raise error when I tried to use Gemini. The only difference my situation with documentation is that I don't have API key in workbench since I use Gemini access from project. Could you help me understand how can I use it? use Gemini config: "model": { "provider": "google", "name": "gemini-pro", # I also tried "gemini-1.0-pro" both raise same error as follows "params": {} },

Error Screenshots

ValidationError Traceback (most recent call last) Cell In[18], line 2 1 # create an aggent for labeling ----> 2 agent = LabelingAgent(config=config)

File /opt/conda/lib/python3.10/site-packages/autolabel/labeler.py:96, in LabelingAgent.init(self, config, cache, example_selector, create_task, console_output, generation_cache, transform_cache) 92 self.config = ( 93 config if isinstance(config, AutolabelConfig) else AutolabelConfig(config) 94 ) 95 self.task = TaskFactory.from_config(self.config) ---> 96 self.llm: BaseModel = ModelFactory.from_config( 97 self.config, cache=self.generation_cache 98 ) 99 score_type = "logprob_average" 100 if self.config.task_type() == TaskType.ATTRIBUTE_EXTRACTION:

File /opt/conda/lib/python3.10/site-packages/autolabel/models/init.py:47, in ModelFactory.from_config(config, cache) 45 try: 46 model_cls = MODEL_REGISTRY[provider] ---> 47 model_obj = model_cls(config=config, cache=cache) 48 # The below ensures that users should based off of the BaseModel 49 # when creating/registering custom models. 50 assert isinstance( 51 model_obj, BaseModel 52 ), f"{model_obj} should inherit from autolabel.models.BaseModel"

File /opt/conda/lib/python3.10/site-packages/autolabel/models/palm.py:68, in PaLMLLM.init(self, config, cache) 66 self.llm = ChatVertexAI(model_name=self.model_name, **self.model_params) 67 else: ---> 68 self.llm = VertexAI(model_name=self.model_name, **self.model_params)

File /opt/conda/lib/python3.10/site-packages/langchain/load/serializable.py:74, in Serializable.init(self, **kwargs) 73 def init(self, **kwargs: Any) -> None: ---> 74 super().init(**kwargs) 75 self._lc_kwargs = kwargs

File /opt/conda/lib/python3.10/site-packages/pydantic/main.py:341, in pydantic.main.BaseModel.init()

ValidationError: 1 validation error for VertexAI root Unknown model publishers/google/models/gemini-1.0-pro; {'gs://google-cloud-aiplatform/schema/predict/instance/text_generation_1.0.0.yaml': <class 'vertexai.preview.language_models._PreviewTextGenerationModel'>} (type=value_error)

Additional context Add any other context about the problem here.

[rishabh-bhargava]

cc: @Vaibhav2001