kandanai commited on
Commit
008c6be
·
verified ·
1 Parent(s): 024b2cf

Update README.md

Browse files
Files changed (1) hide show
  1. README.md +3 -16
README.md CHANGED
@@ -6,20 +6,10 @@
6
  ```
7
  import torch
8
  from transformers import AutoModelForSequenceClassification, AutoTokenizer, BitsAndBytesConfig
9
- from pathlib import Path
10
- import psutil
11
-
12
- def get_memory_usage():
13
- process = psutil.Process()
14
- memory_info = process.memory_info()
15
- return memory_info.rss / 1024**2 # Convert to MB
16
 
17
  def load_and_compute_scores_with_quantized_model(model_path):
18
- model_name = "BAAI/bge-reranker-v2-m3"
19
- tokenizer = AutoTokenizer.from_pretrained(model_name)
20
-
21
- bnb_config = BitsAndBytesConfig(load_in_8bit=True)
22
- model = AutoModelForSequenceClassification.from_pretrained(model_path, config=bnb_config)
23
 
24
  def compute_score(pairs):
25
  inputs = tokenizer(pairs, padding=True, truncation=True, return_tensors='pt')
@@ -27,12 +17,9 @@ def load_and_compute_scores_with_quantized_model(model_path):
27
  outputs = model(**inputs)
28
  return outputs.logits
29
 
30
- after_load_memory = get_memory_usage()
31
- print(f"Memory Usage after loading model: {after_load_memory:.2f} MB")
32
-
33
  scores = compute_score([['what is panda?', 'hi'], ['what is panda?', 'The giant panda (Ailuropoda melanoleuca), sometimes called a panda bear or simply panda, is a bear species endemic to China.']])
34
  print("Scores:", scores)
35
 
36
  quantized_model_path = "quantized_bge_reranker_v2_m3"
37
- load_and_compute_scores_with_quantized_model(Path(quantized_model_path))
38
  ```
 
6
  ```
7
  import torch
8
  from transformers import AutoModelForSequenceClassification, AutoTokenizer, BitsAndBytesConfig
 
 
 
 
 
 
 
9
 
10
  def load_and_compute_scores_with_quantized_model(model_path):
11
+ tokenizer = AutoTokenizer.from_pretrained(model_path)
12
+ model = AutoModelForSequenceClassification.from_pretrained(model_path)
 
 
 
13
 
14
  def compute_score(pairs):
15
  inputs = tokenizer(pairs, padding=True, truncation=True, return_tensors='pt')
 
17
  outputs = model(**inputs)
18
  return outputs.logits
19
 
 
 
 
20
  scores = compute_score([['what is panda?', 'hi'], ['what is panda?', 'The giant panda (Ailuropoda melanoleuca), sometimes called a panda bear or simply panda, is a bear species endemic to China.']])
21
  print("Scores:", scores)
22
 
23
  quantized_model_path = "quantized_bge_reranker_v2_m3"
24
+ load_and_compute_scores_with_quantized_model(quantized_model_path)
25
  ```