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

Update Arch-Guard and corresponding e2e test

Browse source
This commit is contained in:
Shuguang Chen 2024-12-06 13:41:18 -08:00
parent 6dc92fbbc1
commit 2fd8a5a06d
2 changed files with 15 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

6
model_server/app/model_handler/guardrails.py
View file

@ -105,7 +105,7 @@ class ArchGuardHanlder:
sentence = None
return GuardResponse(
prob=prob.item(), verdict=verdict, sentence=sentence, latency=latency
prob=[prob.item()], verdict=verdict, sentence=[sentence], latency=latency
)
def predict(self, req: GuardRequest, max_num_words=300) -> GuardResponse:
@ -138,9 +138,9 @@ class ArchGuardHanlder:
chunk_result = self._predict_text(req.task, chunk)
if chunk_result.verdict:
prob.append(chunk_result.prob)
prob.append(chunk_result.prob[0])
verdict = True
sentence.append(chunk_result.sentence)
sentence.append(chunk_result.sentence[0])
latency += chunk_result.latency
return GuardResponse(