readme update

README.md

@@ -1,22 +1,46 @@
 ---
 license: llama3.1
 language: en
-base_model: meta-llama/Llama-3.1-8B-Instruct
+base_model: meta-llama/Meta-Llama-3.1-8B-Instruct
 ---
 
-# MathBite/self_corrective_llama_3.1_8B_end
+# Self-Corrective Llama 3.1 8B
 
-This is a fine-tuned version of `meta-llama/Llama-3.1-8B-Instruct`. The LoRA adapter has been merged into the base model. This model includes a custom hallucination detection head that can intervene during generation to insert corrective instructions.
+This is a fine-tuned version of `meta-llama/Meta-Llama-3.1-8B-Instruct`, augmented with a novel self-correction mechanism designed to mitigate hallucinations. The LoRA adapter has been merged into the base model for easy deployment.
+
+This model features a custom **hallucination detection head** that works in parallel with the main language model. When it detects a potential error in its own generated text, it can insert special instructions like `[rewrite sentence]` or `[rewrite response]` into the output, effectively flagging its own mistakes for correction. This makes the model more reliable for tasks requiring factual accuracy.
+
+## How it Works
+
+The model, an instance of the custom `SelfCorrectiveLlama` class, adds a small, efficient hallucination detection module to the standard Llama architecture. This module analyzes the model's internal states (hidden states) at each generation step to predict the likelihood of a hallucination.
+
+The model's custom `generate` method then uses these predictions. If a hallucination is likely, it overrides the standard token generation process to insert a corrective instruction. This entire process happens in a single forward pass, making it significantly more efficient than multi-step, agent-based correction pipelines that require multiple LLM calls.
+
+## Intended Use & Prompting
+
+This model is intended for tasks where factual accuracy and faithfulness to a source context are critical, such as question answering or summarization.
+
+While it can be used with standard prompts, its self-correction behavior was reinforced during training using a specific instruction. To achieve the best results and fully leverage the self-correction mechanism, you should include the following note in your system prompt or at the beginning of your input:
+
+<br>
+
+> **Note on Self-Correction**: As you generate your response, you may encounter an automated instruction. This indicates a potential error was detected.
+> - If you see the instruction `[rewrite sentence]`, it means the preceding sentence is incorrect. You must immediately provide a new, corrected version of that sentence.
+> - If you see the instruction `[rewrite response]`, it means the entire preceding response is incorrect. You must immediately provide a new, complete response from the beginning.
+
+<br>
 
 ## How to Use
 
 Because this model uses a custom architecture with a modified `generate` method, you **must** use `trust_remote_code=True` when loading it. The required `modeling.py` file is included in this repository.
 
+**Note:** The custom `generate` method currently only supports a **batch size of 1**.
+
 ```python
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 
-model_name = "MathBite/self_corrective_llama_3.1_8B_end"
+model_name = "MathBite/self_corrective_llama_3.1_8B"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 
 # Important: You must trust the remote code
@@ -26,8 +50,20 @@ model = AutoModelForCausalLM.from_pretrained(
     torch_dtype=torch.bfloat16 # or your preferred dtype
 ).to("cuda") # move model to GPU
 
-# You can now use the model's custom generate method
-prompt = "YOUR PROMPT HERE" # your prompt here
+# Example prompt with the self-correction instruction
+prompt = """
+...
+Note on Self-Correction: As you generate your response, you may encounter an automated instruction. This indicates a potential error was detected.
+- If you see the instruction `[rewrite sentence]`, it means the preceding sentence is incorrect. You must immediately provide a new, corrected version of that sentence.
+- If you see the instruction `[rewrite response]`, it means the entire preceding response is incorrect. You must immediately provide a new, complete response from the beginning.
+
+---
+
+Context: The Eiffel Tower is a wrought-iron lattice tower on the Champ de Mars in Paris, France. It is named after the engineer Gustave Eiffel, whose company designed and built the tower.
+
+Question: Who was the first person to climb the Eiffel Tower?
+"""
+
 inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
 
 # The custom generate method requires the tokenizer instance
@@ -44,6 +80,6 @@ print(generated_text)
 
 ## Model Details
 
-This model was programmatically merged and uploaded using a deployment script. The custom class `SelfCorrectiveLlama` can be found in the `modeling.py` file.
+This model was programmatically merged and uploaded using a deployment script. The custom class `SelfCorrectiveLlama` can be found in the `modeling.py` file included in this repository.
 
-The code in `modeling.py` is licensed under the Apache 2.0 License. The model weights are subject to the original license of the base model.
+The code in `modeling.py` is licensed under the Apache 2.0 License. The model weights are subject to the original license of the base model.