Malkevich Dmitriy
03/05/2024, 2:23 PMimport bentoml
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
@bentoml.service
class Service:
def __init__(self) -> None:
checkpoint = "bigcode/starcoder"
model = AutoModelForCausalLM.from_pretrained(checkpoint)
tokenizer = AutoTokenizer.from_pretrained(checkpoint)
self.pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, device=0)
@bentoml.api
def generate(self, prompt: str):
result = self.pipe(prompt)
return result