Spaces:
Sleeping
Sleeping
Commit
·
c740076
1
Parent(s):
7e7599b
update config kwargs
Browse files
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,
|
| 60 |
-
generator_mini = pipeline(task="text-generation", model=ankerbot_model, tokenizer=ankerbot_tokenizer, torch_dtype=torch.float16,
|
| 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
|