File size: 1,563 Bytes
3933325
 
 
 
 
 
 
 
 
 
 
 
 
 
00b72fa
 
60e4f79
00b72fa
8a76514
 
00b72fa
3933325
 
 
7745907
 
 
 
3933325
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
2103519
7745907
2103519
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
"""Model hosted on Hugging face.

Based on: https://huggingface.co/docs/hub/spaces-sdks-docker-first-demo
"""

from fastapi import FastAPI, Request

from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
from transformers import T5Tokenizer, T5ForConditionalGeneration


# FROM: https://huggingface.co/facebook/blenderbot-400M-distill?text=Hey+my+name+is+Thomas%21+How+are+you%3F
# tokenizer = AutoTokenizer.from_pretrained("facebook/blenderbot-400M-distill")
# model = AutoModelForSeq2SeqLM.from_pretrained("facebook/blenderbot-400M-distill")
# tokenizer = AutoTokenizer.from_pretrained("facebook/blenderbot-1B-distill")
# model = AutoModelForSeq2SeqLM.from_pretrained("facebook/blenderbot-1B-distill")
# tokenizer = T5Tokenizer.from_pretrained("google/flan-t5-small")
# model = T5ForConditionalGeneration.from_pretrained("google/flan-t5-small")
tokenizer = T5Tokenizer.from_pretrained("google/flan-t5-base")
model = T5ForConditionalGeneration.from_pretrained("google/flan-t5-base")


token_size_limit = 128


app = FastAPI()


@app.post('/reply')
async def Reply(req: Request):
    request = await req.json()
    msg = request['msg']
    print(f'MSG: {msg}')

    input_ids = tokenizer(msg, return_tensors='pt').input_ids  # .to('cuda')
    output = model.generate(
        input_ids[:, -token_size_limit:],
        do_sample=True,
        temperature=0.9,
        max_length=100,
    )
    reply = tokenizer.batch_decode(output)[0]
    print(f'REPLY: {reply}')
    return {'reply': reply}


@app.get("/")
def read_root():
    return {"Hello": "World!"}