Update README.md
Browse files
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 |
-
|
19 |
-
|
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(
|
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 |
```
|