RamAnanth1 commited on
Commit
d6d59d4
1 Parent(s): 11b4644

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +4 -2
app.py CHANGED
@@ -4,7 +4,7 @@ import gradio as gr
4
  from gradio.themes.base import Base
5
  from gradio.themes.utils import colors, fonts, sizes
6
  from instruct_pipeline import InstructionTextGenerationPipeline
7
- from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
8
 
9
  import torch
10
 
@@ -17,7 +17,9 @@ theme = gr.themes.Monochrome(
17
  )
18
 
19
  tokenizer = AutoTokenizer.from_pretrained("databricks/dolly-v2-12b", padding_side="left")
20
- model = AutoModelForCausalLM.from_pretrained("databricks/dolly-v2-12b", device_map="auto", load_in_8bit=True)
 
 
21
 
22
  generate_text = InstructionTextGenerationPipeline(model=model, tokenizer=tokenizer)
23
 
 
4
  from gradio.themes.base import Base
5
  from gradio.themes.utils import colors, fonts, sizes
6
  from instruct_pipeline import InstructionTextGenerationPipeline
7
+ from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline, BitsAndBytesConfig
8
 
9
  import torch
10
 
 
17
  )
18
 
19
  tokenizer = AutoTokenizer.from_pretrained("databricks/dolly-v2-12b", padding_side="left")
20
+ quantization_config = BitsAndBytesConfig(load_in_8bit=True,
21
+ llm_int8_threshold=200.0)
22
+ model = AutoModelForCausalLM.from_pretrained("databricks/dolly-v2-12b", device_map="auto", quantization_config=quantization_config)
23
 
24
  generate_text = InstructionTextGenerationPipeline(model=model, tokenizer=tokenizer)
25