asdc commited on
Commit
8b99594
·
verified ·
1 Parent(s): 4fb4767

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +10 -2
app.py CHANGED
@@ -1,5 +1,6 @@
1
  import gradio as gr
2
- from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 
3
  from peft import PeftModel
4
  import torch
5
  import os
@@ -21,7 +22,14 @@ base_model = AutoModelForCausalLM.from_pretrained(
21
  token=API_KEY
22
  )
23
 
24
- model = PeftModel.from_pretrained(base_model, PEFT_ADAPTER, token=API_KEY, load_in_8bit=True)
 
 
 
 
 
 
 
25
 
26
  pipe = pipeline(
27
  "text-generation",
 
1
  import gradio as gr
2
+ from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline, BitsAndBytesConfig
3
+
4
  from peft import PeftModel
5
  import torch
6
  import os
 
22
  token=API_KEY
23
  )
24
 
25
+ nf4_config = BitsAndBytesConfig(
26
+ load_in_4bit=True,
27
+ bnb_4bit_quant_type="nf4",
28
+ bnb_4bit_use_double_quant=True,
29
+ bnb_4bit_compute_dtype=torch.bfloat16
30
+ )
31
+
32
+ model = PeftModel.from_pretrained(base_model, PEFT_ADAPTER, token=API_KEY, quantization_config=nf4_config)
33
 
34
  pipe = pipeline(
35
  "text-generation",