Ankerkraut commited on
Commit
c740076
·
1 Parent(s): 7e7599b

update config kwargs

Browse files
Files changed (1) hide show
  1. app.py +2 -2
app.py CHANGED
@@ -56,8 +56,8 @@ def load_model():
56
  truncation=True,
57
  padding=True, )
58
  ankerbot_model.to("cuda")
59
- generator = pipeline(task="text-generation", model=ankerbot_model, tokenizer=ankerbot_tokenizer, torch_dtype=torch.float16, attn_implementation="eager", use_flash_attention_2=False) # True for flash-attn2 else False
60
- generator_mini = pipeline(task="text-generation", model=ankerbot_model, tokenizer=ankerbot_tokenizer, torch_dtype=torch.float16, attn_implementation="eager", use_flash_attention_2=False) # True for flash-attn2 else False
61
  return (generator, generator_mini)
62
 
63
  _model_cache = None
 
56
  truncation=True,
57
  padding=True, )
58
  ankerbot_model.to("cuda")
59
+ generator = pipeline(task="text-generation", model=ankerbot_model, tokenizer=ankerbot_tokenizer, torch_dtype=torch.float16, trust_remote_code=False) # True for flash-attn2 else False
60
+ generator_mini = pipeline(task="text-generation", model=ankerbot_model, tokenizer=ankerbot_tokenizer, torch_dtype=torch.float16, trust_remote_code=False) # True for flash-attn2 else False
61
  return (generator, generator_mini)
62
 
63
  _model_cache = None