Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 30 additions & 0 deletions reward_model_training/README.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,30 @@
# 📁 File Structure:

reward_model_training/
├── README.md # This guide
├── main.ipynb # 🎯 MAIN NOTEBOOK - Run this!
├── rm_util.py # utility class for training, loading, and calling the reward model
├── reward_data.jsonl # training data used for training the reward model
├── requirements.txt # Python dependencies
└── reward_model # The saved trained reward model directory.


# Train a reward model:

We'll train the base model microsoft/deberta-v3-base into a reward model with training data from reward_data.jsonl. The trained reward model is saved in 'reward_model' dir.

# Evaluate summaries using the reward model:

We'll load the trained reward model and use it to grade the sample summaries. We'll also evaluate with ROUGE and BERTScore.

# 🎯 Run the main.ipynb to conduct the above tasks:

```bash
jupyter notebook main.ipynb
```

# 🔧 Utility Files

rm_util.py - RewardModelUtil class that has the functionality of training a reward model, loading the reward model, and scoring for summarizations provided.

If the 'reward_model' direcotry is present, the training function will not be called; if this directory does not exist, the training function will be executed.
274 changes: 274 additions & 0 deletions reward_model_training/main.ipynb
Original file line number Diff line number Diff line change
@@ -0,0 +1,274 @@
{
"cells": [
{
"cell_type": "markdown",
"id": "aec83135",
"metadata": {},
"source": [
"# Train the reward model:"
]
},
{
"cell_type": "code",
"execution_count": 1,
"id": "a90e6b47",
"metadata": {
"vscode": {
"languageId": "plaintext"
}
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"🟡 Using MPS (Apple Silicon GPU)\n",
"Base model: microsoft/deberta-v3-base\n",
"Training microsoft/deberta-v3-base with reward_data.jsonl\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"Some weights of DebertaV2ForSequenceClassification were not initialized from the model checkpoint at microsoft/deberta-v3-base and are newly initialized: ['classifier.bias', 'classifier.weight', 'pooler.dense.bias', 'pooler.dense.weight']\n",
"You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n",
"/Library/Frameworks/Python.framework/Versions/3.13/lib/python3.13/site-packages/transformers/convert_slow_tokenizer.py:564: UserWarning: The sentencepiece tokenizer that you are converting to a fast tokenizer uses the byte fallback option which is not implemented in the fast tokenizers. In practice this means that the fast version of the tokenizer can produce unknown tokens whereas the sentencepiece version would have converted these unknown tokens into a sequence of byte tokens matching the original piece of text.\n",
" warnings.warn(\n"
]
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "9bd7eb17d28641d58a77a56da173d71f",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Filtering train >1024 tokens: 0%| | 0/10 [00:00<?, ? examples/s]"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"The tokenizer has new PAD/BOS/EOS tokens that differ from the model config and generation config. The model config and generation config were aligned accordingly, being updated with the tokenizer's values. Updated tokens: {'eos_token_id': 2, 'bos_token_id': 1}.\n",
"/Library/Frameworks/Python.framework/Versions/3.13/lib/python3.13/site-packages/torch/utils/data/dataloader.py:684: UserWarning: 'pin_memory' argument is set as true but not supported on MPS now, then device pinned memory won't be used.\n",
" warnings.warn(warn_msg)\n"
]
},
{
"data": {
"text/html": [
"\n",
" <div>\n",
" \n",
" <progress value='6' max='6' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
" [6/6 00:18, Epoch 3/3]\n",
" </div>\n",
" <table border=\"1\" class=\"dataframe\">\n",
" <thead>\n",
" <tr style=\"text-align: left;\">\n",
" <th>Step</th>\n",
" <th>Training Loss</th>\n",
" </tr>\n",
" </thead>\n",
" <tbody>\n",
" </tbody>\n",
"</table><p>"
],
"text/plain": [
"<IPython.core.display.HTML object>"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"/Library/Frameworks/Python.framework/Versions/3.13/lib/python3.13/site-packages/torch/utils/data/dataloader.py:684: UserWarning: 'pin_memory' argument is set as true but not supported on MPS now, then device pinned memory won't be used.\n",
" warnings.warn(warn_msg)\n",
"/Library/Frameworks/Python.framework/Versions/3.13/lib/python3.13/site-packages/torch/utils/data/dataloader.py:684: UserWarning: 'pin_memory' argument is set as true but not supported on MPS now, then device pinned memory won't be used.\n",
" warnings.warn(warn_msg)\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Loading trained reward model...\n",
"✅ trained reward model loaded successfully\n"
]
}
],
"source": [
"import json\n",
"from rm_util import RewardModelUtil\n",
"from evaluate import load\n",
"\n",
"rm = RewardModelUtil()\n",
"rm.train_model('reward_data.jsonl')\n",
"rm.load_model()"
]
},
{
"cell_type": "markdown",
"id": "633707f2",
"metadata": {},
"source": [
"# Sample summary data to be evaluated with ROUGE, BERTScore, and the trained reward model"
]
},
{
"cell_type": "code",
"execution_count": 2,
"id": "c4c24458",
"metadata": {
"vscode": {
"languageId": "plaintext"
}
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"generated_summaries: ['The architectures of LNU-Net and IBU-Net have a down-sampling path for feature extraction and an up-sampling path for precise localization. We use the original U-Net as the basic segmentation approach and compared it with our proposed architectures.', 'Integrating symbolic constraints into deep learning models could make them more robust but this is a time-consuming and challenging task. In this paper we propose AgenticDomiKnowS (ADS) to eliminate this dependency. ADS translates free-form task descriptions into a complete DomiKnowS program and supports optional human-in-the-loop intervention. We show how ADS enables experienced DomiKnowS users and non-users to rapidly construct neuro-symbolic programs, reducing development time from hours to 10-15 minutes.', 'LLMs need continual learning because the knowledge of LLMs quickly becomes outdated as data evolve. Memory-augmented approaches address this by equipping LLMs with a memory bank, that is an external memory module which stores information for future use. However, the memory bank constantly grows in the real-world scenario. In this paper, we propose MBC, a model that compresses the memory bank through a codebook optimization strategy during online adaptation learning.']\n",
"reference_summaries: ['This papar introduces two deep learning approaches called LNU-Net and IBU-Net for Automated Segmentation of Left Ventricle in Cine Cardiac MRI', 'This paper proposes AgenticDomiKnowS (ADS) which eliminates the dependency of integrating symbolic constraints into deep learning models, and shows how ADS enables experienced DomiKnowS users and non-users to rapidly construct neuro-symbolic programs.', 'Memory bank is an external memory module which stores information for LLM to get new knowledge. This paper proposes MBC model that addresses the problem of the memory bank constantly growing, by compressing the memory bank through a codebook optimization strategy during online adaptation learning, and employing Key-Value Low-Rank Adaptation in the attention layers of the LLM.']\n"
]
}
],
"source": [
"sample_summaries = [\n",
" {\n",
" 'original_content': 'Left ventricle (LV) segmentation is critical for clinical quantification and diagnosis of cardiac images. In this work, we propose two novel deep learning architectures called LNU-Net and IBU-Net for left ventricle segmentation from short-axis cine MRI images. LNU-Net is derived from layer normalization (LN) U-Net architecture, while IBU-Net is derived from the instance-batch normalized (IB) U-Net for medical image segmentation. The architectures of LNU-Net and IBU-Net have a down-sampling path for feature extraction and an up-sampling path for precise localization. We use the original U-Net as the basic segmentation approach and compared it with our proposed architectures. Both LNU-Net and IBU-Net have left ventricle segmentation methods: LNU-Net applies layer normalization in each convolutional block, while IBU-Net incorporates instance and batch normalization together in the first convolutional block and passes its result to the next layer. Our method incorporates affine transformations and elastic deformations for image data processing. Our dataset that contains 805 MRI images regarding the left ventricle from 45 patients is used for evaluation. We experimentally evaluate the results of the proposed approaches outperforming the dice coefficient and the average perpendicular distance than other state-of-the-art approaches.',\n",
" 'generated_summary': 'The architectures of LNU-Net and IBU-Net have a down-sampling path for feature extraction and an up-sampling path for precise localization. We use the original U-Net as the basic segmentation approach and compared it with our proposed architectures.',\n",
" 'reference_summary': 'This papar introduces two deep learning approaches called LNU-Net and IBU-Net for Automated Segmentation of Left Ventricle in Cine Cardiac MRI'\n",
" },\n",
" {\n",
" 'original_content': 'Integrating symbolic constraints into deep learning models could make them more robust, interpretable, and data-efficient. Still, it remains a time-consuming and challenging task. Existing frameworks like DomiKnowS help this integration by providing a high-level declarative programming interface, but they still assume the user is proficient with the library\\'s specific syntax. We propose AgenticDomiKnowS (ADS) to eliminate this dependency. ADS translates free-form task descriptions into a complete DomiKnowS program using an agentic workflow that creates and tests each DomiKnowS component separately. The workflow supports optional human-in-the-loop intervention, enabling users familiar with DomiKnowS to refine intermediate outputs. We show how ADS enables experienced DomiKnowS users and non-users to rapidly construct neuro-symbolic programs, reducing development time from hours to 10-15 minutes.',\n",
" 'generated_summary': 'Integrating symbolic constraints into deep learning models could make them more robust but this is a time-consuming and challenging task. In this paper we propose AgenticDomiKnowS (ADS) to eliminate this dependency. ADS translates free-form task descriptions into a complete DomiKnowS program and supports optional human-in-the-loop intervention. We show how ADS enables experienced DomiKnowS users and non-users to rapidly construct neuro-symbolic programs, reducing development time from hours to 10-15 minutes.',\n",
" 'reference_summary': 'This paper proposes AgenticDomiKnowS (ADS) which eliminates the dependency of integrating symbolic constraints into deep learning models, and shows how ADS enables experienced DomiKnowS users and non-users to rapidly construct neuro-symbolic programs.'\n",
" },\n",
" {\n",
" 'original_content': 'Large Language Models (LLMs) have become a mainstay for many everyday applications. However, as data evolve their knowledge quickly becomes outdated. Continual learning aims to update LLMs with new information without erasing previously acquired knowledge. Although methods such as full fine-tuning can incorporate new data, they are computationally expensive and prone to catastrophic forgetting, where prior knowledge is overwritten. Memory-augmented approaches address this by equipping LLMs with a memory bank, that is an external memory module which stores information for future use. However, these methods face a critical limitation, in particular, the memory bank constantly grows in the real-world scenario when large-scale data streams arrive. In this paper, we propose MBC, a model that compresses the memory bank through a codebook optimization strategy during online adaptation learning. To ensure stable learning, we also introduce an online resetting mechanism that prevents codebook collapse. In addition, we employ Key-Value Low-Rank Adaptation in the attention layers of the LLM, enabling efficient utilization of the compressed memory representations. Experiments with benchmark question-answering datasets demonstrate that MBC reduces the memory bank size to 0.3% when compared against the most competitive baseline, while maintaining high retention accuracy during online adaptation learning. Our code is publicly available at https://github.com/Thomkat/MBC.',\n",
" 'generated_summary': 'LLMs need continual learning because the knowledge of LLMs quickly becomes outdated as data evolve. Memory-augmented approaches address this by equipping LLMs with a memory bank, that is an external memory module which stores information for future use. However, the memory bank constantly grows in the real-world scenario. In this paper, we propose MBC, a model that compresses the memory bank through a codebook optimization strategy during online adaptation learning.',\n",
" 'reference_summary': 'Memory bank is an external memory module which stores information for LLM to get new knowledge. This paper proposes MBC model that addresses the problem of the memory bank constantly growing, by compressing the memory bank through a codebook optimization strategy during online adaptation learning, and employing Key-Value Low-Rank Adaptation in the attention layers of the LLM.'\n",
" }\n",
"]\n",
"\n",
"generated_summaries = [entry['generated_summary'] for entry in sample_summaries]\n",
"reference_summaries = [entry['reference_summary'] for entry in sample_summaries]\n",
"\n",
"print(f'generated_summaries: {generated_summaries}')\n",
"print(f'reference_summaries: {reference_summaries}')"
]
},
{
"cell_type": "markdown",
"id": "a37e1998",
"metadata": {},
"source": [
"# ROUGE and BERTScore:"
]
},
{
"cell_type": "code",
"execution_count": 3,
"id": "da09d1fc",
"metadata": {
"vscode": {
"languageId": "plaintext"
}
},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"Some weights of RobertaModel were not initialized from the model checkpoint at roberta-large and are newly initialized: ['pooler.dense.bias', 'pooler.dense.weight']\n",
"You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"ROUGE: {'rouge1': np.float64(0.4496096801440313), 'rouge2': np.float64(0.31125154231618923), 'rougeL': np.float64(0.3524266486098547), 'rougeLsum': np.float64(0.3524266486098547)}\n",
"BERTScore: {'precision': [0.8630497455596924, 0.8990731239318848, 0.9037832021713257], 'recall': [0.8537697792053223, 0.9558082818984985, 0.910071611404419], 'f1': [0.8583846688270569, 0.9265730381011963, 0.9069164991378784], 'hashcode': 'roberta-large_L17_no-idf_version=0.3.12(hug_trans=4.57.1)'}\n"
]
}
],
"source": [
"rouge = load(\"rouge\")\n",
"bertscore = load(\"bertscore\")\n",
"\n",
"results_rouge = rouge.compute(predictions=generated_summaries, references=reference_summaries)\n",
"results_bertscore = bertscore.compute(predictions=generated_summaries, references=reference_summaries, lang=\"en\")\n",
"\n",
"print(\"ROUGE:\", results_rouge)\n",
"print(\"BERTScore:\", results_bertscore)"
]
},
{
"cell_type": "markdown",
"id": "b4287853",
"metadata": {},
"source": [
"# Reward Model Scores on the generated summaries and reference summaries"
]
},
{
"cell_type": "code",
"execution_count": 4,
"id": "bbbe444c",
"metadata": {
"vscode": {
"languageId": "plaintext"
}
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Reward Model Scores: generated summary: -0.16497120261192322; reference summary: -0.1577276885509491\n",
"Reward Model Scores: generated summary: -0.15327253937721252; reference summary: -0.12794075906276703\n",
"Reward Model Scores: generated summary: -0.17355310916900635; reference summary: -0.1665564775466919\n"
]
}
],
"source": [
"for entry in sample_summaries:\n",
" score1, score2 = rm.score_summaries(entry)\n",
" print(f'Reward Model Scores: generated summary: {score1}; reference summary: {score2}')"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.13.9"
}
},
"nbformat": 4,
"nbformat_minor": 5
}
4 changes: 4 additions & 0 deletions reward_model_training/requirements.txt
Original file line number Diff line number Diff line change
@@ -0,0 +1,4 @@
evaluate
nltk
rouge_score
bert_score
Loading