Skip to content
GitLab
Menu
Projects
Groups
Snippets
Loading...
Help
Help
Support
Community forum
Keyboard shortcuts
?
Submit feedback
Contribute to GitLab
Sign in / Register
Toggle navigation
Menu
Open sidebar
change
sglang
Commits
c1930022
"vscode:/vscode.git/clone" did not exist on "cea0e4b9ead83adec919bb21c62360491317d351"
Unverified
Commit
c1930022
authored
Jul 16, 2024
by
Aidan Cooper
Committed by
GitHub
Jul 16, 2024
Browse files
Add support for VertexAI safety settings (#624)
parent
fe3be159
Changes
1
Show whitespace changes
Inline
Side-by-side
Showing
1 changed file
with
8 additions
and
4 deletions
+8
-4
python/sglang/backend/vertexai.py
python/sglang/backend/vertexai.py
+8
-4
No files found.
python/sglang/backend/vertexai.py
View file @
c1930022
import
os
import
warnings
from
typing
import
List
,
Optional
,
Union
import
numpy
as
np
from
typing
import
Optional
from
sglang.backend.base_backend
import
BaseBackend
from
sglang.lang.chat_template
import
get_chat_template
...
...
@@ -16,12 +14,15 @@ try:
GenerativeModel
,
Image
,
)
from
vertexai.generative_models._generative_models
import
SafetySettingsType
except
ImportError
as
e
:
GenerativeModel
=
e
class
VertexAI
(
BaseBackend
):
def
__init__
(
self
,
model_name
):
def
__init__
(
self
,
model_name
,
safety_settings
:
Optional
[
SafetySettingsType
]
=
None
):
super
().
__init__
()
if
isinstance
(
GenerativeModel
,
Exception
):
...
...
@@ -33,6 +34,7 @@ class VertexAI(BaseBackend):
self
.
model_name
=
model_name
self
.
chat_template
=
get_chat_template
(
"default"
)
self
.
safety_settings
=
safety_settings
def
get_chat_template
(
self
):
return
self
.
chat_template
...
...
@@ -54,6 +56,7 @@ class VertexAI(BaseBackend):
ret
=
GenerativeModel
(
self
.
model_name
).
generate_content
(
prompt
,
generation_config
=
GenerationConfig
(
**
sampling_params
.
to_vertexai_kwargs
()),
safety_settings
=
self
.
safety_settings
,
)
comp
=
ret
.
text
...
...
@@ -78,6 +81,7 @@ class VertexAI(BaseBackend):
prompt
,
stream
=
True
,
generation_config
=
GenerationConfig
(
**
sampling_params
.
to_vertexai_kwargs
()),
safety_settings
=
self
.
safety_settings
,
)
for
ret
in
generator
:
yield
ret
.
text
,
{}
...
...
Write
Preview
Markdown
is supported
0%
Try again
or
attach a new file
.
Attach a file
Cancel
You are about to add
0
people
to the discussion. Proceed with caution.
Finish editing this message first!
Cancel
Please
register
or
sign in
to comment