DauroCamilo commited on
Commit
2605bf3
·
verified ·
1 Parent(s): b55152c

Update main.py

Browse files
Files changed (1) hide show
  1. main.py +6 -0
main.py CHANGED
@@ -1,3 +1,9 @@
 
 
 
 
 
 
1
  from fastapi import FastAPI
2
  from pydantic import BaseModel
3
  from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 
1
+ import os
2
+
3
+ os.environ["HF_HOME"] = "/tmp/hf"
4
+ os.environ["HF_DATASETS_CACHE"] = "/tmp/hf/datasets"
5
+ os.environ["HF_METRICS_CACHE"] = "/tmp/hf/metrics"
6
+ os.environ["TRANSFORMERS_CACHE"] = "/tmp/hf/transformers"
7
  from fastapi import FastAPI
8
  from pydantic import BaseModel
9
  from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer