iamkhadke commited on
Commit
2f7c9e1
·
1 Parent(s): dc1a9ca

Update app.py

Browse files

changes added for torch type.

Files changed (1) hide show
  1. app.py +2 -1
app.py CHANGED
@@ -1,7 +1,8 @@
1
  from transformers import AutoModelForCausalLM, AutoTokenizer
2
  from instruct_pipeline import InstructionTextGenerationPipeline
 
3
  tokenizer = AutoTokenizer.from_pretrained("databricks/dolly-v2-3b", padding_side="left")
4
- model = AutoModelForCausalLM.from_pretrained("databricks/dolly-v2-3b")
5
 
6
  generate_text = InstructionTextGenerationPipeline(model=model, tokenizer=tokenizer)
7
 
 
1
  from transformers import AutoModelForCausalLM, AutoTokenizer
2
  from instruct_pipeline import InstructionTextGenerationPipeline
3
+ import torch
4
  tokenizer = AutoTokenizer.from_pretrained("databricks/dolly-v2-3b", padding_side="left")
5
+ model = AutoModelForCausalLM.from_pretrained("databricks/dolly-v2-3b", torch_dtype=torch.bfloat16)
6
 
7
  generate_text = InstructionTextGenerationPipeline(model=model, tokenizer=tokenizer)
8