DUTwangzhijun commited on
Commit
2477ad8
·
1 Parent(s): a71bad8

Update model.py

Browse files
Files changed (1) hide show
  1. model.py +10 -3
model.py CHANGED
@@ -18,7 +18,14 @@ if torch.cuda.is_available():
18
  use_safetensors=False,
19
  )
20
  else:
21
- model = None
 
 
 
 
 
 
 
22
  tokenizer = AutoTokenizer.from_pretrained(model_id)
23
 
24
 
@@ -66,8 +73,8 @@ def run(message: str,
66
  temperature=temperature,
67
  num_beams=1,
68
  )
69
- # t = Thread(target=model.generate, kwargs=generate_kwargs)
70
- # t.start()
71
 
72
  outputs = []
73
  for text in streamer:
 
18
  use_safetensors=False,
19
  )
20
  else:
21
+ model = AutoModelForCausalLM.from_pretrained(
22
+ model_id,
23
+ config=config,
24
+ torch_dtype=torch.float16,
25
+ load_in_4bit=True,
26
+ device_map='auto',
27
+ use_safetensors=False,
28
+ )
29
  tokenizer = AutoTokenizer.from_pretrained(model_id)
30
 
31
 
 
73
  temperature=temperature,
74
  num_beams=1,
75
  )
76
+ t = Thread(target=model.generate, kwargs=generate_kwargs)
77
+ t.start()
78
 
79
  outputs = []
80
  for text in streamer: