apunkt/plano
1
0
Fork 0
mirror of https://github.com/katanemo/plano.git synced 2026-06-29 15:49:40 +02:00
Code Issues Projects Releases Packages Wiki Activity

lint + formating with black (#158)

Browse source 
* lint + formating with black

* add black as pre commit
This commit is contained in:
Co Tran 2024-10-09 11:25:07 -07:00 • committed by GitHub
parent 498e7f9724
commit 5c4a6bc8ff
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
22 changed files with 581 additions and 295 deletions
Download patch file Download diff file Expand all files Collapse all files

9
model_server/app/main.py
View file

@ -7,7 +7,12 @@ from app.load_models import (
get_device,
)
import os
from app.utils import GuardHandler, split_text_into_chunks, load_yaml_config, get_model_server_logger
from app.utils import (
GuardHandler,
split_text_into_chunks,
load_yaml_config,
get_model_server_logger,
)
import torch
import yaml
import string
@ -39,6 +44,7 @@ guard_handler = GuardHandler(toxic_model=None, jailbreak_model=jailbreak_model)
app = FastAPI()
class EmbeddingRequest(BaseModel):
input: str
model: str
@ -84,6 +90,7 @@ async def embedding(req: EmbeddingRequest, res: Response):
}
return {"data": data, "model": req.model, "object": "list", "usage": usage}
class GuardRequest(BaseModel):
input: str
task: str