Tonic commited on
Commit
38fedf1
1 Parent(s): aae125b

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +4 -3
app.py CHANGED
@@ -1,13 +1,14 @@
1
  import spaces
2
  import gradio as gr
3
  import transformers
4
- from transformers import AutoTokenizer, AutoModelForCausalLM
5
  import torch
6
 
7
- # Load the tokenizer and model
 
8
  model_id = "CohereForAI/c4ai-command-r-v01"
9
  tokenizer = AutoTokenizer.from_pretrained(model_id)
10
- model = AutoModelForCausalLM.from_pretrained(model_id)
11
 
12
  @spaces.GPU
13
  def generate_response(user_input, max_new_tokens, temperature):
 
1
  import spaces
2
  import gradio as gr
3
  import transformers
4
+ from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
5
  import torch
6
 
7
+ bnb_config = BitsAndBytesConfig(load_in_8bit=True)
8
+
9
  model_id = "CohereForAI/c4ai-command-r-v01"
10
  tokenizer = AutoTokenizer.from_pretrained(model_id)
11
+ model = AutoModelForCausalLM.from_pretrained(model_id, quantization_config=bnb_config)
12
 
13
  @spaces.GPU
14
  def generate_response(user_input, max_new_tokens, temperature):