File size: 1,935 Bytes
35d9624 3e86372 5b941ac 3e86372 35d9624 3e86372 eb4b3b5 3e86372 eb4b3b5 3e86372 35d9624 eb4b3b5 35d9624 |
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 |
import time
import json
from pydantic import BaseModel
import transformers
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
import torch
from torch import nn
import torch.nn.functional as F
from torch.cuda.amp import custom_fwd, custom_bwd
from bitsandbytes.functional import quantize_blockwise, dequantize_blockwise
from loguru import logger
from typing import Dict, List, Any
# -----------------------------------------> API <---------------------------------------
name="Kanpredict/gptj-6b-8bits"
model = AutoModelForCausalLM.from_pretrained(name, device_map="auto", load_in_8bit=True)
tokenizer = AutoTokenizer.from_pretrained(name)
class EndpointHandler:
def __init__(self, path=""):
# create inference pipeline
self.pipeline = pipeline(model=name, model_kwargs= {"device_map": "auto", "load_in_8bit": True}, max_new_tokens=max_new_tokens)
def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
inputs = data.pop("inputs", data)
parameters = data.pop("parameters", None)
# run the model and get the output(generated text)
prompt = inputs
temperature = float(parameters.temperature)
length = int(parameters.length)
logger.info("message input: %s", prompt)
logger.info("tempereture: %s", parameters.temperature)
logger.info("length: %s", parameters.length)
start = time.time()
prompt = tokenizer(prompt, return_tensors='pt')
prompt = {key: value.to(device) for key, value in prompt.items()}
out = self.pipeline(**prompt, min_length=length, max_length=length, temperature=temperature, do_sample=True)
generated_text = tokenizer.decode(out[0])
logger.info("generated text: ", generated_text)
logger.info("time taken: %s", time.time() - start)
result = {"output": generated_text}
result = json.dumps(result)
return result
|