CreitinGameplays commited on
Commit
42753e5
1 Parent(s): e36f792

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +10 -1
app.py CHANGED
@@ -1,13 +1,22 @@
1
  import gradio as gr
2
  import torch
 
3
  from transformers import AutoTokenizer, AutoModelForCausalLM
4
 
 
 
 
 
 
 
 
 
5
  # Define the BLOOM model name
6
  model_name = "CreitinGameplays/bloom-3b-conversational"
7
 
8
  # Load tokenizer and model
9
  tokenizer = AutoTokenizer.from_pretrained(model_name)
10
- model = AutoModelForCausalLM.from_pretrained(model_name)
11
 
12
  def generate_text(user_prompt):
13
  """Generates text using the BLOOM model from Hugging Face Transformers and removes the user prompt."""
 
1
  import gradio as gr
2
  import torch
3
+ import bitsandbytes as bnb
4
  from transformers import AutoTokenizer, AutoModelForCausalLM
5
 
6
+ # BNB config
7
+ bnb_config = BitsAndBytesConfig(
8
+ load_in_4bit=True,
9
+ bnb_4bit_use_double_quant=True,
10
+ bnb_4bit_quant_type="nf4",
11
+ bnb_4bit_compute_dtype=torch.bfloat16
12
+ )
13
+
14
  # Define the BLOOM model name
15
  model_name = "CreitinGameplays/bloom-3b-conversational"
16
 
17
  # Load tokenizer and model
18
  tokenizer = AutoTokenizer.from_pretrained(model_name)
19
+ model = AutoModelForCausalLM.from_pretrained(model_name, quantization_config=bnb_config)
20
 
21
  def generate_text(user_prompt):
22
  """Generates text using the BLOOM model from Hugging Face Transformers and removes the user prompt."""