Danil commited on
Commit
98fde76
Β·
1 Parent(s): eba2192
Files changed (1) hide show
  1. server.py +0 -17
server.py DELETED
@@ -1,17 +0,0 @@
1
- import uvicorn
2
- from fastapi import FastAPI
3
- from transformers import AutoModelForCausalLM, AutoTokenizer
4
-
5
- app = FastAPI()
6
-
7
- model_name = 'facebook/incoder-1B'
8
- tokenizer = AutoTokenizer.from_pretrained(model_name)
9
- model = AutoModelForCausalLM.from_pretrained(model_name, low_cpu_mem_usage=True)
10
- print('load ok')
11
-
12
- @app.get("/")
13
- def read_root(input_text, max_length, top_p, top_k, num_beams, temperature, repetition_penalty):
14
- inpt = tokenizer.encode(input_text, return_tensors="pt")
15
- out = model.generate(inpt, max_length=int(max_length), top_p=float(top_p), top_k=float(top_k), temperature=float(temperature), num_beams=int(num_beams), repetition_penalty=float(repetition_penalty))
16
- res = tokenizer.decode(out[0])
17
- return {"text": res}