elplaguister 4n3mone commited on
Commit
9c32b43
β€’
1 Parent(s): 126a8c6

Update koalpaca.py (#8)

Browse files

- Update koalpaca.py (ee299e59aa54e845319082bc51ca2e04b70cc231)


Co-authored-by: yongsang yoo <[email protected]>

Files changed (1) hide show
  1. koalpaca.py +2 -1
koalpaca.py CHANGED
@@ -8,6 +8,7 @@ class KoAlpaca(Model):
8
  def __init__(self):
9
  peft_model_id = "4n3mone/Komuchat-koalpaca-polyglot-12.8B"
10
  config = PeftConfig.from_pretrained(peft_model_id)
 
11
  self.bnb_config = BitsAndBytesConfig(
12
  load_in_4bit=True,
13
  bnb_4bit_use_double_quant=True,
@@ -29,7 +30,7 @@ class KoAlpaca(Model):
29
  inputs,
30
  return_tensors='pt',
31
  return_token_type_ids=False
32
- ).to('cpu'),
33
  generation_config=self.gen_config
34
  )
35
  outputs = self.tokenizer.decode(output_ids[0]).split("### λ‹΅λ³€: ")[-1]
 
8
  def __init__(self):
9
  peft_model_id = "4n3mone/Komuchat-koalpaca-polyglot-12.8B"
10
  config = PeftConfig.from_pretrained(peft_model_id)
11
+ accelerator = Accelerator()
12
  self.bnb_config = BitsAndBytesConfig(
13
  load_in_4bit=True,
14
  bnb_4bit_use_double_quant=True,
 
30
  inputs,
31
  return_tensors='pt',
32
  return_token_type_ids=False
33
+ ).to(accelerator.device),
34
  generation_config=self.gen_config
35
  )
36
  outputs = self.tokenizer.decode(output_ids[0]).split("### λ‹΅λ³€: ")[-1]