husseinelsaadi commited on
Commit
722e882
·
1 Parent(s): 2e9d967
backend/services/resume_parser.py CHANGED
@@ -6,7 +6,6 @@ import zipfile
6
  import json
7
  import torch
8
  from typing import List
9
- from transformers import AutoModelForCausalLM, AutoTokenizer
10
 
11
  os.environ["OMP_NUM_THREADS"] = "1"
12
  os.environ["OPENBLAS_NUM_THREADS"] = "1"
@@ -14,18 +13,25 @@ os.environ["MKL_NUM_THREADS"] = "1"
14
  os.environ["NUMEXPR_NUM_THREADS"] = "1"
15
  os.environ["VECLIB_MAXIMUM_THREADS"] = "1"
16
 
17
- MODEL_ID = "tiiuae/falcon-7b-instruct"
 
18
 
 
 
 
 
 
 
19
 
20
- print(f"Loading {MODEL_ID}... (This may take some time on first run)")
21
- tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
22
  model = AutoModelForCausalLM.from_pretrained(
23
- MODEL_ID,
24
- torch_dtype=torch.float32,
25
  device_map="auto"
26
  )
27
 
28
 
 
29
  # ===============================
30
  # Text Extraction (PDF/DOCX)
31
  # ===============================
 
6
  import json
7
  import torch
8
  from typing import List
 
9
 
10
  os.environ["OMP_NUM_THREADS"] = "1"
11
  os.environ["OPENBLAS_NUM_THREADS"] = "1"
 
13
  os.environ["NUMEXPR_NUM_THREADS"] = "1"
14
  os.environ["VECLIB_MAXIMUM_THREADS"] = "1"
15
 
16
+ from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
17
+ import torch
18
 
19
+ bnb_config = BitsAndBytesConfig(
20
+ load_in_4bit=True,
21
+ bnb_4bit_compute_dtype=torch.float16,
22
+ bnb_4bit_use_double_quant=True,
23
+ bnb_4bit_quant_type="nf4"
24
+ )
25
 
26
+ tokenizer = AutoTokenizer.from_pretrained("tiiuae/falcon-7b-instruct")
 
27
  model = AutoModelForCausalLM.from_pretrained(
28
+ "tiiuae/falcon-7b-instruct",
29
+ quantization_config=bnb_config,
30
  device_map="auto"
31
  )
32
 
33
 
34
+
35
  # ===============================
36
  # Text Extraction (PDF/DOCX)
37
  # ===============================
requirements.txt CHANGED
@@ -61,4 +61,5 @@ requests>=2.31.0
61
  # Spaces. The version is pinned for reproducibility and compatibility.
62
  psycopg2-binary
63
  matplotlib
 
64
 
 
61
  # Spaces. The version is pinned for reproducibility and compatibility.
62
  psycopg2-binary
63
  matplotlib
64
+ bitsandbytes>=0.41.0
65