File size: 1,858 Bytes
997becd 229a1e9 997becd 3555efd 229a1e9 997becd 229a1e9 80ef686 229a1e9 997becd 80ef686 997becd 80ef686 997becd 229a1e9 997becd 229a1e9 997becd 229a1e9 997becd 229a1e9 997becd |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 |
import torch
from typing import Dict, List, Any
from transformers import LlamaForCausalLM, LlamaTokenizer, pipeline
from pynvml import nvmlInit, nvmlDeviceGetHandleByIndex, nvmlDeviceGetMemoryInfo
nvmlInit()
gpu_h1 = nvmlDeviceGetHandleByIndex(0)
print('loaded_imports')
# get dtype
dtype = torch.bfloat16 if torch.cuda.get_device_capability()[0] == 8 else torch.float16
print('chose dtype', dtype)
class EndpointHandler:
def __init__(self, path=""):
# load the model
print('starting to load tokenizer')
tokenizer = LlamaTokenizer.from_pretrained("/repository/orca_tokenizer", local_files_only=True)
print('loaded tokenizer')
gpu_info1 = nvmlDeviceGetMemoryInfo(gpu_h1)
print(f'vram {gpu_info1.total} used {gpu_info1.used} free {gpu_info1.free}')
model = LlamaForCausalLM.from_pretrained(
"/repository/pytorch_model",
device_map="auto",
torch_dtype=dtype,
offload_folder="offload",
local_files_only=True
)
gpu_info1 = nvmlDeviceGetMemoryInfo(gpu_h1)
print(f'vram {gpu_info1.total} used {gpu_info1.used} free {gpu_info1.free}')
print('loaded model')
# create inference pipeline
self.pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer)
print('created pipeline')
def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
print('starting to call')
inputs = data.pop("inputs", data)
print('inputs: ', inputs)
parameters = data.pop("parameters", None)
# pass inputs with all kwargs in data
if parameters is not None:
prediction = self.pipeline(inputs, **parameters)
else:
prediction = self.pipeline(inputs)
# postprocess the prediction
return prediction
|