Update README.md

README.md

@@ -79,30 +79,57 @@ Python
 
 import torch
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
+from peft import PeftModel
+from huggingface_hub import snapshot_download # Import the downloader tool
 
-# Define the model ID and the reward model subfolder
-model_id = "nabeelshan/rlhf-gpt2-pipeline"
-subfolder = "reward_model_final"
+# --- CONFIGURATION ---
+BASE_MODEL_ID = "openai-community/gpt2"
+HF_MODEL_ID = "nabeelshan/rlhf-gpt2-pipeline"
+SUBFOLDER = "reward_model_final"
+
+print(f"Downloading model files from '{HF_MODEL_ID}'...")
+local_model_path = snapshot_download(
+    repo_id=HF_MODEL_ID,
+    allow_patterns=f"{SUBFOLDER}/*"
+)
+local_adapter_path = f"{local_model_path}/{SUBFOLDER}"
+print(f"   Successfully downloaded to: {local_adapter_path}")
+
+
+print("Loading model from local path...")
+tokenizer = AutoTokenizer.from_pretrained(local_adapter_path)
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+
+base_model = AutoModelForSequenceClassification.from_pretrained(
+    BASE_MODEL_ID,
+    num_labels=1,
+    pad_token_id=tokenizer.pad_token_id
+)
+
+model = PeftModel.from_pretrained(base_model, local_adapter_path)
+model.eval()
+print("   Model loaded successfully!")
 
-# Load the tokenizer and reward model
-tokenizer = AutoTokenizer.from_pretrained(model_id, subfolder=subfolder)
-model = AutoModelForSequenceClassification.from_pretrained(model_id, subfolder=subfolder)
 
 prompt = "What diet should I follow to lose weight healthily?"
 good_response = "A balanced, nutritious plan based on eating whole foods is best. Limit processed and sugary foods."
 bad_response = "Just eat less lol."
 
-# Tokenize the inputs (prompt + response)
-inputs_good = tokenizer(prompt, good_response, return_tensors="pt")
-inputs_bad = tokenizer(prompt, bad_response, return_tensors="pt")
+def get_reward_score(prompt_text: str, response_text: str) -> float:
+    """Tokenizes and calculates the reward score for a given prompt and response."""
+    inputs = tokenizer(prompt_text, response_text, return_tensors="pt", padding=True, truncation=True)
+    with torch.no_grad():
+        result = model(**inputs)
+        return result.logits[0].item()
+
+score_good = get_reward_score(prompt, good_response)
+score_bad = get_reward_score(prompt, bad_response)
+
+print(f"\nScore for good response: {score_good:.2f}")
+print(f"Score for bad response:  {score_bad:.2f}")
 
-# Get the reward scores (logits)
-with torch.no_grad():
-    reward_good = model(**inputs_good).logits[0].item()
-    reward_bad = model(**inputs_bad).logits[0].item()
 
-print(f"Score for good response: {reward_good:.2f}")
-print(f"Score for bad response: {reward_bad:.2f}")
 
 # The model should give a higher score to the better response.
 # Expected: Score for good response: 2.15