README.md

---
license: mit
tags:
- ppo
- qlora
- reinforcement-learning
- llama-3
- mmlu
pipeline_tag: text-generation
---

# PPO-QLoRA Trained Model (spark-model-QLoRA)

This repository contains an agent (actor and critic models) trained using Proximal Policy Optimization (PPO) with QLoRA.
The training was performed using the scripts and models available in the `spark_rl` directory of the `explore-rl` project.

**Base Model:** `meta-llama/Llama-3-8B-Instruct` (or specify if different, based on your `train.py` arguments)

## Model Components

The `model_final` directory (uploaded here as the root of these components) contains:

*   **`actor/`**: LoRA adapters for the actor (policy) model.
*   **`critic/`**: LoRA adapters for the critic (value) model's base LLM, and a `value_head.pt` file for its custom value prediction head.
*   **`tokenizer/`**: The Hugging Face tokenizer used during training.
*   **`hyperparams.txt`**: Key hyperparameters used for the PPO training.
*   **`models.py`**: Contains the `LLMActorLora` and `LLMCriticLora` class definitions required to load and use these models.

## How to Use

To use these models, you will need the `LLMActorLora` and `LLMCriticLora` classes from the included `models.py` file.

```python
import torch
from transformers import AutoTokenizer
from models import LLMActorLora, LLMCriticLora # models.py is in this repository

# --- Configuration ---
BASE_MODEL_ID = "meta-llama/Llama-3-8B-Instruct" # IMPORTANT: Ensure this matches the model used for training!
MODEL_REPO_PATH = "gabrielbo/spark-model-QLoRA" # Or local path if downloaded
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# --- Load Tokenizer ---
try:
    tokenizer = AutoTokenizer.from_pretrained(f"{MODEL_REPO_PATH}/tokenizer")
except Exception: # Fallback if tokenizer is in the root
    tokenizer = AutoTokenizer.from_pretrained(MODEL_REPO_PATH)

if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "left" # Ensure consistency if PPO agent used left padding

# --- Load Actor ---
actor = LLMActorLora(
    device=DEVICE,
    model_id=BASE_MODEL_ID,
    # lora_r and disable_quantization can be defaults or from hyperparams.txt
)
# Path to actor adapters within the model repo
actor_adapters_path = f"{MODEL_REPO_PATH}/actor" 
actor.load_pretrained(actor_adapters_path)
actor.model.eval()
print("Actor loaded successfully.")

# --- Load Critic ---
critic = LLMCriticLora(
    device=DEVICE,
    model_id=BASE_MODEL_ID,
    # lora_r and disable_quantization can be defaults or from hyperparams.txt
)
# Path to critic components within the model repo
critic_components_path = f"{MODEL_REPO_PATH}/critic"
critic.load_pretrained(critic_components_path)
critic.model.eval()
critic.value_head.eval()
print("Critic loaded successfully.")

# --- Example: Generating an action (conceptual) ---
# This part is highly dependent on how your PPOAgent prepares inputs.
# The following is a generic example. You'll need to adapt it.

# Example input construction (refer to PPOAgent.prepare_batch)
question = "What is the capital of France?"
state_text = "The current context is a geography quiz."
input_text = f"Question: {question}\n\nState: {state_text}\n\nAction:"

inputs = tokenizer(input_text, return_tensors="pt", padding=True, truncation=True, max_length=512).to(DEVICE)

print(f"\nGenerating action for: {input_text}")
with torch.no_grad():
    # Actor generates token IDs
    # Note: Generation kwargs might be needed (e.g., temperature, top_p from hyperparams.txt or evaluate.py)
    generated_ids = actor.generate(
        inputs.input_ids,
        attention_mask=inputs.attention_mask,
        max_new_tokens=50, # Adjust as needed
        # temperature=0.7, # Example
        # top_p=0.9,       # Example
        do_sample=True   # Example, if sampling was used
    )
    
    # Decode the generated action
    # The generated output includes the input_text, so we need to slice it off.
    # This depends on tokenizer.padding_side; if "left", then slicing logic changes.
    # Assuming tokenizer.padding_side = "right" (default for many models) or handled by generate
    
    # If tokenizer.padding_side was "left" for generation, the input is at the end.
    # For simplicity, let's assume the output only contains new tokens after input.
    # This might need adjustment based on specific generation config.
    
    # A common way to get only the generated part:
    response_ids = generated_ids[0][inputs.input_ids.shape[-1]:]
    action_text = tokenizer.decode(response_ids, skip_special_tokens=True)
    
    print(f"Generated Action: {action_text.strip()}")

    # --- Example: Getting a value estimate (conceptual) ---
    value_prediction = critic.forward(inputs.input_ids, attention_mask=inputs.attention_mask)
    print(f"Value prediction for the state: {value_prediction.item()}")

```

## Training Details

The model was trained using the PPO algorithm with the following key settings (see `hyperparams.txt` for more details):

*   **Learning Rate (Actor)**: (Refer to `lr` in `hyperparams.txt`)
*   **Learning Rate (Critic)**: (Refer to `critic_lr` in `hyperparams.txt`)
*   **PPO Clip Ratio**: (Refer to `clip_ratio` in `hyperparams.txt`)
*   **KL Coefficient**: (Refer to `kl_coef` in `hyperparams.txt`)
*   **Target KL**: (Refer to `target_kl` in `hyperparams.txt`)
*   **Batch Size**: (As per your training script, e.g., `args.batch`)
*   **PPO Epochs**: (As per your training script, e.g., `args.ppo_epochs`)
*   **Total PPO Iterations**: (As per your training script, e.g., `args.steps`)

The specific dataset used for training was MMLU trajectories.

## Intended Use

This model is intended for tasks requiring sequential decision-making and reasoning, similar to the MMLU benchmark. It can be used as a starting point for further fine-tuning or for direct application in relevant domains.

## Limitations

*   The model's performance is tied to the quality and characteristics of the offline trajectory data it was trained on.
*   As a LoRA-adapted model, it relies on the capabilities of the base `meta-llama/Llama-3-8B-Instruct` model.
*   The generation behavior may require careful prompt engineering.

## Citation

If you use this model or the `spark_rl` codebase, please consider citing the original `explore-rl` repository:
[Link to your explore-rl GitHub repository, if public]