snsynth commited on
Commit
46f6ce5
·
1 Parent(s): 8980288

add hf variables

Browse files
Files changed (1) hide show
  1. app.py +5 -1
app.py CHANGED
@@ -1,3 +1,7 @@
 
 
 
 
1
  from fastapi import FastAPI
2
  from pydantic import BaseModel
3
  from transformers import AutoTokenizer, AutoModel
@@ -22,7 +26,7 @@ class EmbeddingRequest(BaseModel):
22
  def root():
23
  return {"message": "API is working"}
24
 
25
-
26
  @app.post("/v1/embeddings")
27
  def create_embeddings(request: EmbeddingRequest):
28
  with torch.no_grad():
 
1
+ import os
2
+ os.environ["HF_HOME"] = "/app/hf_home"
3
+ os.environ["TRANSFORMERS_CACHE"] = "/app/hf_home/transformers"
4
+
5
  from fastapi import FastAPI
6
  from pydantic import BaseModel
7
  from transformers import AutoTokenizer, AutoModel
 
26
  def root():
27
  return {"message": "API is working"}
28
 
29
+
30
  @app.post("/v1/embeddings")
31
  def create_embeddings(request: EmbeddingRequest):
32
  with torch.no_grad():