tmmdev commited on
Commit
f4ebb84
·
verified ·
1 Parent(s): 5055efe

Update pattern_analyzer.py

Browse files
Files changed (1) hide show
  1. pattern_analyzer.py +10 -6
pattern_analyzer.py CHANGED
@@ -1,14 +1,20 @@
1
  import os
2
  os.environ['HF_HOME'] = '/tmp/huggingface'
3
 
4
- from transformers import AutoModelForCausalLM, AutoTokenizer
5
  import torch
6
  import json
7
  import pandas as pd
8
  from pattern_logic import PatternLogic
9
 
 
 
10
  class PatternAnalyzer:
11
  def __init__(self):
 
 
 
 
 
12
  model_kwargs = {
13
  "device_map": "auto",
14
  "torch_dtype": torch.float32,
@@ -18,18 +24,16 @@ class PatternAnalyzer:
18
  "disk": "8GB"
19
  },
20
  "offload_folder": "/tmp/offload",
21
- "load_in_8bit": True,
22
- "revision": "main"
23
  }
24
 
25
  self.model = AutoModelForCausalLM.from_pretrained(
26
  "tmmdev/codellama-pattern-analysis",
27
  **model_kwargs,
28
- use_safetensors=True,
29
- trust_remote_code=True,
30
- resume_download=True
31
  )
32
 
 
33
  self.tokenizer = AutoTokenizer.from_pretrained(
34
  "tmmdev/codellama-pattern-analysis",
35
  use_fast=True
 
1
  import os
2
  os.environ['HF_HOME'] = '/tmp/huggingface'
3
 
 
4
  import torch
5
  import json
6
  import pandas as pd
7
  from pattern_logic import PatternLogic
8
 
9
+ from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
10
+
11
  class PatternAnalyzer:
12
  def __init__(self):
13
+ quantization_config = BitsAndBytesConfig(
14
+ load_in_8bit=True,
15
+ llm_int8_threshold=6.0
16
+ )
17
+
18
  model_kwargs = {
19
  "device_map": "auto",
20
  "torch_dtype": torch.float32,
 
24
  "disk": "8GB"
25
  },
26
  "offload_folder": "/tmp/offload",
27
+ "quantization_config": quantization_config
 
28
  }
29
 
30
  self.model = AutoModelForCausalLM.from_pretrained(
31
  "tmmdev/codellama-pattern-analysis",
32
  **model_kwargs,
33
+ trust_remote_code=True
 
 
34
  )
35
 
36
+
37
  self.tokenizer = AutoTokenizer.from_pretrained(
38
  "tmmdev/codellama-pattern-analysis",
39
  use_fast=True