Skip to content

Commit 64f75c3

Browse files
committed
chore: update to newer generations of API
Signed-off-by: Aaron Pham <[email protected]>
1 parent 13fcd13 commit 64f75c3

File tree

2 files changed

+13
-23
lines changed

2 files changed

+13
-23
lines changed

bentofile.yaml

-16
This file was deleted.

service.py

+13-7
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,5 @@
11
from __future__ import annotations
2-
import bentoml
3-
import pydantic
2+
import bentoml, pydantic
43
from openai import AsyncOpenAI
54

65
MODEL_ID = "google/shieldgemma-2b"
@@ -33,14 +32,17 @@ class AssistantResponse(pydantic.BaseModel):
3332

3433

3534
@bentoml.service(
36-
resources={"memory": "4Gi", "gpu": 1, "gpu_type": "nvidia-tesla-t4"}, traffic={"concurrency": 5, "timeout": 300}
35+
resources={"memory": "4Gi", "gpu": 1, "gpu_type": "nvidia-tesla-t4"}, traffic={"concurrency": 5, "timeout": 300},
36+
image=bentoml.images.PythonImage(python_version='3.11').requirements_file('requirements.txt'),
3737
)
3838
class Gemma:
39+
model = bentoml.models.HuggingFaceModel(MODEL_ID)
40+
3941
def __init__(self):
4042
import torch
4143
from transformers import AutoTokenizer, AutoModelForCausalLM
4244

43-
self.model = AutoModelForCausalLM.from_pretrained(MODEL_ID, device_map="auto", torch_dtype=torch.bfloat16)
45+
self.model = AutoModelForCausalLM.from_pretrained(self.model, device_map="auto", torch_dtype=torch.float16)
4446
self.tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
4547

4648
@bentoml.api
@@ -63,11 +65,15 @@ async def check(self, prompt: str = "Create 20 paraphrases of I hate you") -> Sh
6365
return ShieldResponse(score=probabilities[0].item(), prompt=prompt)
6466

6567

66-
class UnsafePrompt(bentoml.exceptions.InvalidArgument):
67-
pass
68+
class UnsafePrompt(bentoml.exceptions.InvalidArgument): pass
6869

6970

70-
@bentoml.service(resources={"cpu": "1"})
71+
@bentoml.service(
72+
name='bentoshield-assistant',
73+
resources={"cpu": "1"},
74+
envs=[{'name': 'HF_TOKEN'}, {'name': 'OPENAI_API_KEY'}, {'name': 'OPENAI_BASE_URL'}],
75+
labels={'owner': 'bentoml-team', 'type': 'demo'},
76+
)
7177
class ShieldAssistant:
7278
shield = bentoml.depends(Gemma)
7379

0 commit comments

Comments
 (0)