Elaineyy commited on
Commit
e35ad58
Β·
verified Β·
1 Parent(s): 6d22f86

Update server.py

Browse files
Files changed (1) hide show
  1. server.py +2 -10
server.py CHANGED
@@ -6,18 +6,10 @@ import torch
6
 
7
  app = FastAPI()
8
 
9
- # βœ… Fix: Set writable cache directory for Hugging Face models
10
- os.environ["TRANSFORMERS_CACHE"] = "/tmp"
11
- os.environ["HF_HOME"] = "/tmp"
12
-
13
- # βœ… Ensure cache directory exists
14
- if not os.path.exists("/tmp"):
15
- os.makedirs("/tmp")
16
-
17
  # βœ… Load DeepSeek-Coder-V2-Base Model with `trust_remote_code=True`
18
  model_name = "deepseek-ai/DeepSeek-Coder-V2-Base"
19
- tokenizer = AutoTokenizer.from_pretrained(model_name, cache_dir="/tmp", trust_remote_code=True)
20
- model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map="auto", cache_dir="/tmp", trust_remote_code=True)
21
 
22
  class CodeRequest(BaseModel):
23
  user_story: str
 
6
 
7
  app = FastAPI()
8
 
 
 
 
 
 
 
 
 
9
  # βœ… Load DeepSeek-Coder-V2-Base Model with `trust_remote_code=True`
10
  model_name = "deepseek-ai/DeepSeek-Coder-V2-Base"
11
+ tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
12
+ model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map="auto", trust_remote_code=True)
13
 
14
  class CodeRequest(BaseModel):
15
  user_story: str