Daniel Fried commited on
Commit
5b5750d
1 Parent(s): 488ca72

switch to big model

Browse files
Files changed (1) hide show
  1. modules/app.py +11 -4
modules/app.py CHANGED
@@ -17,11 +17,18 @@ import json
17
  # })
18
  # origins=[f"http://localhost:{PORT}", "https://huggingface.co", "https://hf.space"]
19
 
20
- #CUDA = True
21
- CUDA = False
22
  PORT = 7860
23
  VERBOSE = False
24
 
 
 
 
 
 
 
 
 
 
25
  from fastapi import FastAPI, Request
26
  from fastapi.staticfiles import StaticFiles
27
  from fastapi.responses import FileResponse, StreamingResponse
@@ -30,9 +37,9 @@ app.mount("/static", StaticFiles(directory="static"), name="static")
30
 
31
 
32
  print("loading model")
33
- model = AutoModelForCausalLM.from_pretrained("facebook/incoder-1B")
34
  print("loading tokenizer")
35
- tokenizer = AutoTokenizer.from_pretrained("facebook/incoder-1B")
36
  print("loading complete")
37
 
38
  if CUDA:
 
17
  # })
18
  # origins=[f"http://localhost:{PORT}", "https://huggingface.co", "https://hf.space"]
19
 
 
 
20
  PORT = 7860
21
  VERBOSE = False
22
 
23
+ BIG_MODEL = True
24
+
25
+ if big_model:
26
+ CUDA = True
27
+ model_name = "facebook/incoder-6B"
28
+ else:
29
+ CUDA = False
30
+ model_name = "facebook/incoder-1B"
31
+
32
  from fastapi import FastAPI, Request
33
  from fastapi.staticfiles import StaticFiles
34
  from fastapi.responses import FileResponse, StreamingResponse
 
37
 
38
 
39
  print("loading model")
40
+ model = AutoModelForCausalLM.from_pretrained(model_name)
41
  print("loading tokenizer")
42
+ tokenizer = AutoTokenizer.from_pretrained(model_name)
43
  print("loading complete")
44
 
45
  if CUDA: