base for the application : support custom python application and auto-install dependencies

This commit is contained in:
faraphel 2025-01-08 19:00:47 +01:00
parent a20371e1ab
commit e2ebbf8a82
25 changed files with 494 additions and 2 deletions

View file

@ -0,0 +1,11 @@
{
"type": "python",
"file": "model.py",
"requirements": [
"transformers",
"torch",
"torchvision",
"torchaudio"
]
}

View file

@ -0,0 +1,28 @@
import json
import torch
import transformers
MODEL_NAME: str = "huawei-noah/TinyBERT_General_4L_312D"
def load(model):
model.model = transformers.AutoModel.from_pretrained(MODEL_NAME)
model.tokenizer = transformers.AutoTokenizer.from_pretrained(MODEL_NAME)
def unload(model):
model.model = None
model.tokenizer = None
def infer(model, payload: dict) -> str:
inputs = model.tokenizer(payload["prompt"], return_tensors="pt")
with torch.no_grad():
outputs = model.model(**inputs)
embeddings = outputs.last_hidden_state
return json.dumps({
"data": embeddings.tolist()
})