service-internal commited on
Commit
8cf8359
·
verified ·
1 Parent(s): 1e1cbba

Update main.py

Browse files
Files changed (1) hide show
  1. main.py +7 -2
main.py CHANGED
@@ -1,15 +1,20 @@
 
 
 
 
1
  from fastapi import FastAPI, Request
2
  from pydantic import BaseModel
3
  from transformers import pipeline
4
 
 
 
 
5
  # Load zero-shot classification pipeline
6
  classifier = pipeline(
7
  "zero-shot-classification",
8
  model="MoritzLaurer/deberta-v3-large-zeroshot-v2.0"
9
  )
10
 
11
- # Define FastAPI app
12
- app = FastAPI()
13
 
14
  # Define input schema
15
  class InputText(BaseModel):
 
1
+ import os
2
+ os.environ["TRANSFORMERS_CACHE"] = "/tmp/hf-cache"
3
+ os.environ["HF_HOME"] = "/tmp/hf-home"
4
+
5
  from fastapi import FastAPI, Request
6
  from pydantic import BaseModel
7
  from transformers import pipeline
8
 
9
+ # Define FastAPI app
10
+ app = FastAPI()
11
+
12
  # Load zero-shot classification pipeline
13
  classifier = pipeline(
14
  "zero-shot-classification",
15
  model="MoritzLaurer/deberta-v3-large-zeroshot-v2.0"
16
  )
17
 
 
 
18
 
19
  # Define input schema
20
  class InputText(BaseModel):