chayanbhansali commited on
Commit
bb99597
·
verified ·
1 Parent(s): fc3343b

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +5 -5
app.py CHANGED
@@ -11,13 +11,13 @@ class RAGChatbot:
11
  embedding_model="all-MiniLM-L6-v2"):
12
  # Initialize tokenizer and model
13
  self.tokenizer = AutoTokenizer.from_pretrained(model_name)
14
- self.bnb_config = BitsAndBytesConfig(
15
- load_in_8bit=True, # Enable 8-bit loading
16
- llm_int8_threshold=6.0, # Threshold for mixed-precision computation
17
- )
18
  self.model = AutoModelForCausalLM.from_pretrained(
19
  model_name,
20
- quantization_config= self.bnb_config,
21
  device_map="auto"
22
  )
23
 
 
11
  embedding_model="all-MiniLM-L6-v2"):
12
  # Initialize tokenizer and model
13
  self.tokenizer = AutoTokenizer.from_pretrained(model_name)
14
+ # self.bnb_config = BitsAndBytesConfig(
15
+ # load_in_8bit=True, # Enable 8-bit loading
16
+ # llm_int8_threshold=6.0, # Threshold for mixed-precision computation
17
+ # )
18
  self.model = AutoModelForCausalLM.from_pretrained(
19
  model_name,
20
+ torch_dtype=torch.bfloat16, ,
21
  device_map="auto"
22
  )
23