Spaces:
Runtime error
Runtime error
File size: 2,544 Bytes
26a9031 1f85c85 26a9031 1f85c85 26a9031 1f85c85 26a9031 1f85c85 26a9031 |
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 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 |
from fastapi import FastAPI
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse
import torch
# from transformers import pipeline
from transformers import AutoTokenizer, AutoModelForCausalLM
app = FastAPI()
# MODEL = "google/flan-t5-small"
# MODEL = "jingyaogong/minimind-v1-small"
MODEL = "tclh123/minimind-v1-small"
# pipe_flan = pipeline("text2text-generation", model=MODEL, trust_remote_code=True)
device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
tokenizer = AutoTokenizer.from_pretrained(MODEL, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(MODEL, trust_remote_code=True)
model = model.to(device)
model = model.eval()
def query(message, max_seq_len=512, temperature=0.7, top_k=16):
prompt = '请问，' + message
messages = []
messages.append({"role": "user", "content": prompt})
stream = True
# print(messages)
new_prompt = tokenizer.apply_chat_template(
messages,
tokenize=False,
add_generation_prompt=True
)[-(max_seq_len - 1):]
x = tokenizer(new_prompt).data['input_ids']
x = (torch.tensor(x, dtype=torch.long, device=device)[None, ...])
res_y = model.generate(x, tokenizer.eos_token_id, max_new_tokens=max_seq_len, temperature=temperature, top_k=top_k, stream=stream)
try:
y = next(res_y)
except StopIteration:
# print("No answer")
return ""
ret = []
history_idx = 0
while y != None:
answer = tokenizer.decode(y[0].tolist())
if answer and answer[-1] == '�':
try:
y = next(res_y)
except:
break
continue
# print(answer)
if not len(answer):
try:
y = next(res_y)
except:
break
continue
# print(answer[history_idx:], end='', flush=True)
ret.append(answer[history_idx:])
try:
y = next(res_y)
except:
break
history_idx = len(answer)
if not stream:
break
# print('\n')
ret.append('\n')
return ''.join(ret)
@app.get("/infer_t5")
def t5(input):
# output = pipe_flan(input)
# return {"output": output[0]["generated_text"]}
output = query(input)
return {"output": output}
app.mount("/", StaticFiles(directory="static", html=True), name="static")
@app.get("/")
def index() -> FileResponse:
return FileResponse(path="/app/static/index.html", media_type="text/html")
|