minhdang commited on
Commit
4e4d4fa
1 Parent(s): 959022f

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +6 -1
app.py CHANGED
@@ -5,7 +5,12 @@ import gradio as gr
5
  import spaces
6
  import torch
7
  from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 
8
 
 
 
 
 
9
  MAX_MAX_NEW_TOKENS = 2048
10
  DEFAULT_MAX_NEW_TOKENS = 1024
11
  total_count=0
@@ -15,7 +20,7 @@ MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096"))
15
  DESCRIPTION="""CODE"""
16
 
17
  model_id = "deepseek-ai/deepseek-coder-7b-instruct-v1.5"
18
- model = AutoModelForCausalLM.from_pretrained(model_id,device_map="auto",load_in_4bit=True)
19
  tokenizer=AutoTokenizer.from_pretrained(model_id)
20
  tokenizer.use_defaul_system_prompt=True
21
 
 
5
  import spaces
6
  import torch
7
  from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
8
+ from transformers import BitsAndBytesConfig
9
 
10
+ nf4_config = BitsAndBytesConfig(
11
+ load_in_4bit=True,
12
+ bnb_4bit_quant_type="nf4",
13
+ )
14
  MAX_MAX_NEW_TOKENS = 2048
15
  DEFAULT_MAX_NEW_TOKENS = 1024
16
  total_count=0
 
20
  DESCRIPTION="""CODE"""
21
 
22
  model_id = "deepseek-ai/deepseek-coder-7b-instruct-v1.5"
23
+ model = AutoModelForCausalLM.from_pretrained(model_id,device_map="auto",quantization_config=nf4_config)
24
  tokenizer=AutoTokenizer.from_pretrained(model_id)
25
  tokenizer.use_defaul_system_prompt=True
26