Adapters
English
code
dnnsdunca commited on
Commit
4260961
1 Parent(s): 7b93fce

Create App.py

Browse files
Files changed (1) hide show
  1. App.py +22 -0
App.py ADDED
@@ -0,0 +1,22 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from fastapi import FastAPI
2
+ from pydantic import BaseModel
3
+ from transformers import GPT2LMHeadModel, GPT2Tokenizer
4
+ import uvicorn
5
+
6
+ class CodeRequest(BaseModel):
7
+ prompt: str
8
+
9
+ app = FastAPI()
10
+
11
+ model = GPT2LMHeadModel.from_pretrained('./codegen_model')
12
+ tokenizer = GPT2Tokenizer.from_pretrained('./codegen_model')
13
+
14
+ @app.post("/generate-code/")
15
+ def generate_code(request: CodeRequest):
16
+ inputs = tokenizer.encode(request.prompt, return_tensors='pt')
17
+ outputs = model.generate(inputs, max_length=150, num_return_sequences=1)
18
+ generated_code = tokenizer.decode(outputs[0], skip_special_tokens=True)
19
+ return {"generated_code": generated_code}
20
+
21
+ if __name__ == "__main__":
22
+ uvicorn.run(app, host="0.0.0.0", port=8000)