win2win commited on
Commit
523127e
·
verified ·
1 Parent(s): 132a52f

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +28 -0
app.py ADDED
@@ -0,0 +1,28 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from fastapi import FastAPI
2
+ from transformers import pipeline
3
+ import torch
4
+ from pydantic import BaseModel
5
+
6
+ app = FastAPI()
7
+
8
+ # Disable xformers if needed
9
+ torch.backends.cuda.enable_flash_sdp(False)
10
+ torch.backends.cuda.enable_mem_efficient_sdp(False)
11
+
12
+ # Load model
13
+ model = pipeline(
14
+ "text-classification",
15
+ model="win2win/3-epochs-classifier-ver2",
16
+ device="cuda" if torch.cuda.is_available() else "cpu"
17
+ )
18
+
19
+ class Request(BaseModel):
20
+ text: str
21
+
22
+ @app.post("/predict")
23
+ async def predict(request: Request):
24
+ return model(request.text)
25
+
26
+ if __name__ == "__main__":
27
+ import uvicorn
28
+ uvicorn.run(app, host="0.0.0.0", port=8000)