Spaces:
Runtime error
Runtime error
File size: 966 Bytes
9bf2007 dcd2d54 3e6fc0f 9bf2007 5ed2b9f e48a0c0 e5e2748 21e7dd1 fdc39d2 3e6fc0f 6254e11 72231f4 6254e11 72231f4 21e7dd1 5bc3efc 21e7dd1 72231f4 21e7dd1 |
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 |
import transformers
import torch
from fastapi import FastAPI
from transformers import AutoModelForCausalLM, AutoTokenizer
app = FastAPI()
MODEL = None
TOKENIZER = None
# ?input=%22Name%203%20shows%22
@app.get("/")
def llama(input):
prompt = [{'role': 'user', 'content': input}]
inputs = TOKENIZER.apply_chat_template( prompt, add_generation_prompt=True, return_tensors='pt' )
tokens = MODEL.generate( inputs.to(MODEL.device), max_new_tokens=1024, temperature=0.3, do_sample=True)
tresponse = TOKENIZER.decode(tokens[0], skip_special_tokens=False)
print(tresponse)
return tresponse
@app.on_event("startup")
def init_model():
global MODEL
global TOKENIZER
if not MODEL:
print("loading model")
TOKENIZER = AutoTokenizer.from_pretrained('stabilityai/stablelm-zephyr-3b')
MODEL = AutoModelForCausalLM.from_pretrained('stabilityai/stablelm-zephyr-3b', device_map="auto")
print("loaded model") |