Update README.md

README.md

@@ -19,12 +19,74 @@ It has been trained using [TRL](https://github.com/huggingface/trl).
 ## Quick start
 
 ```python
-from transformers import pipeline
+import torch
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+from peft import PeftModel, PeftConfig
+
+# 1. Define the PEFT model ID
+peft_model_id = "xxccho/margin_reg_baseline"
+
+# 2. Load the PEFT config
+config = PeftConfig.from_pretrained(peft_model_id)
+
+# 3. Load tokenizer from base model (safer)
+tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
+
+# Llama padding fix
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+
+# 4. Load base model
+base_model = AutoModelForSequenceClassification.from_pretrained(
+    config.base_model_name_or_path,
+    num_labels=1,
+    torch_dtype=torch.bfloat16,
+    device_map="auto"
+)
+
+# 5. Apply LoRA adapter
+model = PeftModel.from_pretrained(base_model, peft_model_id)
+model.config.pad_token_id = tokenizer.pad_token_id
+model.eval()
+
+# -----------------------------
+# Example Usage (chat format)
+# -----------------------------
+messages = [
+    {"role": "user", "content": "What is the capital of France?"},
+    {"role": "assistant", "content": "The capital of France is Paris."}
+]
+
+# Format prompt using chat template
+formatted_prompt = tokenizer.apply_chat_template(
+    messages,
+    tokenize=False,
+    add_generation_prompt=False
+)
+
+inputs = tokenizer(formatted_prompt, return_tensors="pt").to(model.device)
+
+# Get reward score
+with torch.no_grad():
+    outputs = model(**inputs)
+    reward_score = outputs.logits.squeeze().item()
+
+print(f"[Chat] Reward Score: {reward_score:.4f}")
+
+
+# -----------------------------
+# Example Usage (plain text)
+# -----------------------------
+text = "User: What is the capital of France?\nAssistant: Paris."
+
+inputs = tokenizer(text, return_tensors="pt").to(model.device)
+
+with torch.no_grad():
+    outputs = model(**inputs)
+    reward_score = outputs.logits.squeeze().item()
+
+print(f"[Plain] Reward Score: {reward_score:.4f}")
 
-text = "The capital of France is Paris."
-rewarder = pipeline(model="None", device="cuda")
-output = rewarder(text)[0]
-print(output["score"])
 ```
 
 ## Training procedure