File size: 1,369 Bytes
b166186 1a91849 b166186 88e7970 b166186 88e7970 b166186 |
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 |
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
import os
import subprocess
# Manually install bitsandbytes
def install(package):
subprocess.check_call([sys.executable, "-m", "pip", "install", package])
try:
import bitsandbytes
except ImportError:
install("bitsandbytes==0.39.1")
try:
import accelerate
except ImportError:
install("accelerate==0.20.0")
class ModelHandler:
def __init__(self):
self.model = None
self.tokenizer = None
def load_model(self):
# Load token as env var
model_id = "NiCETmtm/llama3_torch"
token = os.getenv("HF_API_TOKEN")
# Load model & tokenizer
self.model = AutoModelForCausalLM.from_pretrained(model_id, use_auth_token=token, trust_remote_code=True, from_tf=True)
self.tokenizer = AutoTokenizer.from_pretrained(model_id, use_auth_token=token, trust_remote_code=True)
def predict(self, inputs):
tokens = self.tokenizer(inputs, return_tensors="pt")
with torch.no_grad():
outputs = self.model.generate(**tokens)
return self.tokenizer.decode(outputs[0], skip_special_tokens=True)
model_handler = ModelHandler()
model_handler.load_model()
def inference(event, context):
inputs = event["data"]
outputs = model_handler.predict(inputs)
return {"predictions": outputs}
|