ffreemt commited on
Commit
cfea67f
1 Parent(s): 5b46843

Update for CPU: float(), requried

Browse files
Files changed (1) hide show
  1. app.py +9 -6
app.py CHANGED
@@ -27,8 +27,8 @@ except Exception:
27
  # Windows
28
  logger.warning("Windows, cant run time.tzset()")
29
 
30
- # model_name = "THUDM/chatglm2-6b"
31
- model_name = "THUDM/chatglm2-6b-int4"
32
 
33
  RETRY_FLAG = False
34
 
@@ -43,13 +43,16 @@ has_cuda = torch.cuda.is_available()
43
  # has_cuda = False # force cpu
44
 
45
  if has_cuda:
46
- model = (
47
- AutoModel.from_pretrained(model_name, trust_remote_code=True).cuda().half()
48
- ) # 3.92G
 
 
 
49
  else:
50
  model = AutoModel.from_pretrained(
51
  model_name, trust_remote_code=True
52
- ).half() # .float() .half().float()
53
 
54
  model = model.eval()
55
 
 
27
  # Windows
28
  logger.warning("Windows, cant run time.tzset()")
29
 
30
+ # model_name = "THUDM/chatglm2-6b" # 7x?G
31
+ model_name = "THUDM/chatglm2-6b-int4" # 3.9G
32
 
33
  RETRY_FLAG = False
34
 
 
43
  # has_cuda = False # force cpu
44
 
45
  if has_cuda:
46
+ if model_name.endswith("int4"):
47
+ model = AutoModel.from_pretrained(model_name, trust_remote_code=True).cuda()
48
+ else:
49
+ model = (
50
+ AutoModel.from_pretrained(model_name, trust_remote_code=True).cuda().half()
51
+ )
52
  else:
53
  model = AutoModel.from_pretrained(
54
  model_name, trust_remote_code=True
55
+ ).float() # .half().float(), .float() required for CPU
56
 
57
  model = model.eval()
58