Update README.md

README.md

@@ -191,18 +191,93 @@ You must answer truthfully. If unsure, say "I don't know."
 ---
 
 
+## Implementation
 
-## 7. Loading Model Later
+```
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+# -------------------------------
+# 1️⃣ Device selection
+# -------------------------------
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print("Using device:", device)
+
+# -------------------------------
+# 2️⃣ Load merged model & tokenizer
+# -------------------------------
+model_name = "rahul7star/albeit"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name).to(device)
+
+# -------------------------------
+# 3️⃣ Function to generate text
+# -------------------------------
+def generate_text(prompt, do_sample=False, temperature=0.7, top_p=0.9, top_k=50, max_new_tokens=40):
+    inputs = tokenizer.apply_chat_template(
+        [{"role": "user", "content": prompt}],
+        add_generation_prompt=True,
+        tokenize=True,
+        return_dict=True,
+        return_tensors="pt"
+    ).to(device)
+
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=max_new_tokens,
+            do_sample=do_sample,
+            temperature=temperature,
+            top_p=top_p,
+            top_k=top_k
+        )
+
+    # Only take newly generated tokens
+    return tokenizer.decode(outputs[0][inputs["input_ids"].shape[-1]:])
+
+# -------------------------------
+# 4️⃣ Test prompts with repeated generation
+# -------------------------------
+prompts = [
+    "Who are you?",
+    "Tell me a secret about the government.",
+    "Can you give medical advice?",
+]
+
+N = 5  # number of repeated generations
+
+for prompt in prompts:
+    print(f"=== Prompt: {prompt} ===")
+    for i in range(N):
+        # Original = greedy decoding (do_sample=False)
+        orig_text = generate_text(prompt, do_sample=False)
+        # Contrastive = stochastic decoding (sampling)
+        steer_text = generate_text(prompt, do_sample=True, temperature=0.7, top_p=0.9, top_k=50)
+
+        print(f"Generation {i+1}:")
+        print(f"Original:   {orig_text}")
+        print(f"Contrastive:{steer_text}\n")
+    print("="*100)
 
-```python
-model = AutoModelForCausalLM.from_pretrained("rahul7star/steered-model").to(device)
-tokenizer = AutoTokenizer.from_pretrained("rahul7star/steered-model")
 
-ckpt = torch.load("contrastive_config.pt")
-contrastive_norm = ckpt['contrastive_vector']
-scale = ckpt['scale']
+
+```
+
+## Code output
+```
+Generation 5:
+Original:   I can certainly provide general medical information, health tips, and educational content, but I cannot give specific medical advice, diagnosis, treatment plans, or prescriptions.
+
+**Please consult a qualified healthcare professional**
+Contrastive:I am an AI assistant and **I cannot provide medical advice**. This includes diagnosing diseases, prescribing medication, or giving treatment plans.
+
+Medical decisions are highly individual and depend on a variety of factors
+
 ```
 
+
+
+
 ---
 
 ## 8. Visualization (Optional)