Update VertexAIComponent configuration

This commit is contained in:
anovazzi1 2024-01-10 09:34:13 -03:00
commit a291564b4e

View file

@ -11,45 +11,49 @@ class VertexAIComponent(CustomComponent):
return { return {
"credentials": { "credentials": {
"display_name": "Credentials", "display_name": "Credentials",
"type": "file", "field_type": "file",
"file_types": ["json"], "file_types": ["json"],
"required": False, "required": False,
"default": None, "value": None,
}, },
"location": { "location": {
"display_name": "Location", "display_name": "Location",
"type": "str", "type": "str",
"default": "us-central1", "advanced":True,
"value": "us-central1",
"required": False, "required": False,
}, },
"max_output_tokens": { "max_output_tokens": {
"display_name": "Max Output Tokens", "display_name": "Max Output Tokens",
"type": "int", "field_type": "int",
"default": 128, "value": 128,
"required": False, "required": False,
"advanced":True
}, },
"max_retries": { "max_retries": {
"display_name": "Max Retries", "display_name": "Max Retries",
"type": "int", "type": "int",
"default": 6, "value": 6,
"required": False, "required": False,
"advanced":True
}, },
"metadata": { "metadata": {
"display_name": "Metadata", "display_name": "Metadata",
"type": "dict", "field_type": "dict",
"required": False, "required": False,
"default": {}, "default": {},
}, },
"model_name": { "model_name": {
"display_name": "Model Name", "display_name": "Model Name",
"type": "str", "type": "str",
"default": "text-bison", "value": "text-bison",
"required": False, "required": False,
}, },
"n": { "n": {
"advanced":True,
"display_name": "N", "display_name": "N",
"type": "int", "field_type": "int",
"default": 1, "value": 1,
"required": False, "required": False,
}, },
"project": { "project": {
@ -60,57 +64,72 @@ class VertexAIComponent(CustomComponent):
}, },
"request_parallelism": { "request_parallelism": {
"display_name": "Request Parallelism", "display_name": "Request Parallelism",
"type": "int", "field_type": "int",
"default": 5, "value": 5,
"required": False, "required": False,
"advanced":True
}, },
"streaming": { "streaming": {
"display_name": "Streaming", "display_name": "Streaming",
"type": "bool", "field_type": "bool",
"default": False, "value": False,
"required": False, "required": False,
"advanced":True
}, },
"temperature": { "temperature": {
"display_name": "Temperature", "display_name": "Temperature",
"type": "float", "field_type": "float",
"default": 0.0, "value": 0.0,
"required": False, "required": False,
"advanced":True
}, },
"top_k": { "top_k": {
"display_name": "Top K", "display_name": "Top K",
"type": "int", "type": "int",
"default": 40, "default": 40,
"required": False, "required": False,
"advanced":True
}, },
"top_p": { "top_p": {
"display_name": "Top P", "display_name": "Top P",
"type": "float", "field_type": "float",
"default": 0.95, "value": 0.95,
"required": False, "required": False,
"advanced":True
}, },
"tuned_model_name": { "tuned_model_name": {
"display_name": "Tuned Model Name", "display_name": "Tuned Model Name",
"type": "str", "type": "str",
"required": False, "required": False,
"default": None, "value": None,
"advanced":True
}, },
"verbose": { "verbose": {
"display_name": "Verbose", "display_name": "Verbose",
"type": "bool", "field_type": "bool",
"default": False, "value": False,
"required": False, "required": False,
}, },
"name":{
"display_name":"Name",
"field_type":"str"
},
"client_preview":{
"display_name":"client_preview"
}
} }
def build( def build(
self, self,
credentials: Optional[str] = None, credentials: Optional[str] = None,
client_preview: Optional[any]= None,
location: str = "us-central1", location: str = "us-central1",
max_output_tokens: int = 128, max_output_tokens: int = 128,
max_retries: int = 6, max_retries: int = 6,
metadata: Dict = None, metadata: Dict = None,
model_name: str = "text-bison", model_name: str = "text-bison",
n: int = 1, n: int = 1,
name:Optional[str] = None,
project: Optional[str] = None, project: Optional[str] = None,
request_parallelism: int = 5, request_parallelism: int = 5,
streaming: bool = False, streaming: bool = False,
@ -127,6 +146,7 @@ class VertexAIComponent(CustomComponent):
from langchain.llms import VertexAI from langchain.llms import VertexAI
return VertexAI( return VertexAI(
client_preview=client_preview,
credentials=credentials, credentials=credentials,
location=location, location=location,
max_output_tokens=max_output_tokens, max_output_tokens=max_output_tokens,
@ -134,6 +154,7 @@ class VertexAIComponent(CustomComponent):
metadata=metadata, metadata=metadata,
model_name=model_name, model_name=model_name,
n=n, n=n,
name=name,
project=project, project=project,
request_parallelism=request_parallelism, request_parallelism=request_parallelism,
streaming=streaming, streaming=streaming,