Update README.md

README.md

@@ -42,6 +42,112 @@ There was decent amount of empty assistant response on original dataset. So, I d
 
 ---
 
+## Example Code
+
+**The code below is modified from** (**PairRM-hf Repo**)[https://huggingface.co/llm-blender/PairRM-hf]
+
+```python
+import os
+os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+from llm_blender.pair_ranker.pairrm import DebertaV2PairRM
+from transformers import AutoTokenizer
+from typing import List
+pairrm = DebertaV2PairRM.from_pretrained("maywell/Better-PairRM", device_map="cuda:0").eval()
+tokenizer = AutoTokenizer.from_pretrained("maywell/Better-PairRM")
+source_prefix = "<|source|>"
+cand1_prefix = "<|candidate1|>"
+cand2_prefix = "<|candidate2|>"
+inputs = ["hello!", "I love you!"]
+candidates_A = ["hi!", "I hate you!"]
+candidates_B = ["f**k off!", "I love you, too!"]
+def tokenize_pair(sources:List[str], candidate1s:List[str], candidate2s:List[str], source_max_length=2030, candidate_max_length=670):
+    ids = []
+    assert len(sources) == len(candidate1s) == len(candidate2s)
+    max_length = source_max_length + 2 * candidate_max_length
+    for i in range(len(sources)):
+        source_ids = tokenizer.encode(source_prefix + sources[i], max_length=source_max_length, truncation=True)
+        candidate_max_length = (max_length - len(source_ids)) // 2
+        candidate1_ids = tokenizer.encode(cand1_prefix + candidate1s[i], max_length=candidate_max_length, truncation=True)
+        candidate2_ids = tokenizer.encode(cand2_prefix + candidate2s[i], max_length=candidate_max_length, truncation=True)
+        ids.append(source_ids + candidate1_ids + candidate2_ids)
+    encodings = tokenizer.pad({"input_ids": ids}, return_tensors="pt", padding="max_length", max_length=max_length)
+    return encodings
+
+encodings = tokenize_pair(inputs, candidates_A, candidates_B)
+encodings = {k:v.to(pairrm.device) for k,v in encodings.items()}
+outputs = pairrm(**encodings)
+logits = outputs.logits.tolist()
+comparison_results = outputs.logits > 0
+print(logits)
+print(comparison_results)
+```
+
+You can also easily compare two conversations like the followings:
+```python
+from transformers import AutoTokenizer
+
+tokenizer = AutoTokenizer.from_pretrained("microsoft/deberta-v3-large")
+
+def truncate_texts(text, max_length, truncate_side):
+    tokenizer.truncation_side = truncate_side
+    tokens = tokenizer.encode(text, add_special_tokens=False, max_length=max_length)
+    truncated_text = tokenizer.decode(tokens, skip_special_tokens=True)
+    return truncated_text
+
+MY_JINJA_TEMPLATE = """{% for message in messages -%}
+{% if message['role'] == 'user' -%}
+USER: {{ message['content']|trim -}}
+{% if not loop.last -%}
+
+
+{% endif %}
+{% elif message['role'] == 'assistant' -%}
+ASSISTANT: {{ message['content']|trim -}}
+{% if not loop.last -%}
+
+
+{% endif %}
+{% elif message['role'] == 'user_context' -%}
+USER: {{ message['content']|trim -}}
+{% if not loop.last -%}
+
+
+{% endif %}
+{% elif message['role'] == 'system' -%}
+SYSTEM MESSAGE: {{ message['content']|trim -}}
+{% if not loop.last -%}
+
+
+{% endif %}
+{% endif %}
+{% endfor -%}
+{% if add_generation_prompt and messages[-1]['role'] != 'assistant' -%}
+ASSISTANT: {% endif -%}"""
+
+my_jinja2_env = jinja2.Environment()
+my_jinja2_template = my_jinja2_env.from_string(MY_JINJA_TEMPLATE)
+
+def tokenize_conv_pair(convAs: List[str], convBs: List[str]):
+
+    # check conversations correctness
+    assert len(convAs) == len(convBs), "Number of conversations must be the same"
+    for c_a, c_b in zip(convAs, convBs):
+        assert len(c_a) == len(c_b), "Number of turns in each conversation must be the same"
+        assert all([c_a[i]['content'] == c_b[i]['content'] for i in range(0, len(c_a), 2)]), "USER turns must be the same"
+    
+    inputs = [
+        truncate_texts(my_jinja2_template.render(messages=x[:-1], add_generation_prompt=True), 2030, "left") for x in convAs
+    ]
+    cand1_texts = [
+        truncate_texts(x[-1]['content'], 670, "right") for x in convAs
+    ]
+    cand2_texts = [
+        truncate_texts(x[-1]['content'], 670, "right") for x in convBs
+    ]
+    encodings = tokenize_pair(inputs, cand1_texts, cand2_texts)
+    return encodings
+```
+
 ## Statistics
 
 ### Context length