ariG23498 HF staff commited on
Commit
d46c5c3
·
verified ·
1 Parent(s): ccff382

Create main.py

Browse files
Files changed (1) hide show
  1. main.py +16 -0
main.py ADDED
@@ -0,0 +1,16 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from fastapi import FastAPI, Query
2
+ from transformers import pipeline
3
+
4
+ app = FastAPI()
5
+
6
+ # Load a different model, replace "bert-base-uncased" with your desired model
7
+ pipe = pipeline(model="openai-community/gpt2")
8
+
9
+ @app.get("/")
10
+ def read_root():
11
+ return {"message": "API is live. Use the /predict endpoint."}
12
+
13
+ @app.get("/predict")
14
+ def predict(text: str = Query(..., description="Input text for model inference")):
15
+ result = pipe(text)
16
+ return {"predictions": result}