pleonova commited on
Commit
8bb4440
1 Parent(s): 635dc92

Update for predict

Browse files
Files changed (1) hide show
  1. app.py +7 -3
app.py CHANGED
@@ -1,13 +1,17 @@
1
- import os
2
  from fastapi import FastAPI
3
  from transformers import pipeline
 
4
 
5
- # Set a custom Hugging Face cache directory
6
- os.environ["TRANSFORMERS_CACHE"] = "/app/.cache"
7
 
8
  app = FastAPI()
9
  classifier = pipeline("zero-shot-classification", model="facebook/bart-large-mnli")
10
 
 
 
 
 
11
  @app.post("/predict")
12
  async def predict(data: dict):
13
  labels = ["Mathematics", "Language Arts", "Social Studies", "Science"]
 
 
1
  from fastapi import FastAPI
2
  from transformers import pipeline
3
+ import os
4
 
5
+ # Set Hugging Face cache directory
6
+ os.environ["HF_HOME"] = "/app/.cache"
7
 
8
  app = FastAPI()
9
  classifier = pipeline("zero-shot-classification", model="facebook/bart-large-mnli")
10
 
11
+ @app.get("/")
12
+ async def read_root():
13
+ return {"message": "Welcome to the Zero-Shot Classification API"}
14
+
15
  @app.post("/predict")
16
  async def predict(data: dict):
17
  labels = ["Mathematics", "Language Arts", "Social Studies", "Science"]