Stop using pickle, please!

I noticed that we are currently relying on Python's built-in `pickle` module (or standard PyTorch `.pt`/`.pth` file saving, which uses `pickle` under the hood) to store and load our machine learning models.

While `pickle` is incredibly convenient and a historical default in the ML ecosystem, it introduces major architectural and security issues that we should look into addressing. I'd love to propose migrating our weight serialization to Hugging Face's **`safetensors`** format.

#### Why `pickle` is a bad idea for ML models:

1. **Arbitrary Code Execution (RCE):** The `pickle` format is fundamentally a stack of execution instructions (opcodes). When someone calls `pickle.load()`, Python actually executes code embedded inside the file to reconstruct objects. A bad actor can easily package a reverse shell or malware inside a model file, which executes automatically upon load.
2. **Brittle across dependency updates:** `pickle` saves object definitions rather than pure raw data. If we restructure our codebase, rename a module, or update dependencies, older pickle files will completely break because Python won't know how to reconstruct the saved classes.
3. **No Zero-Copy Loading:** Unpickling forces a lot of memory allocations and CPU overhead.

#### Community Resources & Proofs of Danger:

* **Trail of Bits** published an amazing breakdown on how attackers exploit pickle files via [Sleepy Pickle attacks](https://blog.trailofbits.com/2024/06/11/exploiting-ml-models-with-pickle-file-attacks-part-1/) to inject backdoors stealthily.
* **Hugging Face's Security Documentation** on why they are actively moving away from it: [Hugging Face Pickle Scanning & Risks](https://huggingface.co/docs/hub/security-pickle).
* **Snyk Labs** wrote a comprehensive review of these architectural flaws in deep learning: [Vulnerabilities in Deep Learning File Formats](https://labs.snyk.io/resources/vulnerabilities-in-deep-learning-file-formats/).

---

#### How we can migrate to `safetensors`

`safetensors` only saves raw tensor byte buffers and a safe, text-based JSON header for metadata. This completely prevents any hidden code execution.

It’s straightforward to implement since it plays perfectly with PyTorch and NumPy. Here is how simple a migration would look:

**1. Install the library:**

```bash
pip install safetensors

```

**2. Updating our saving logic:**
Instead of `torch.save(model.state_dict(), "model.pt")`, we extract the state dict or standard tensors and pass it directly:

```python
from safetensors.torch import save_file

# Saving PyTorch tensors securely
save_file(model.state_dict(), "model.safetensors")

```

**3. Updating our loading logic:**
Instead of `torch.load()`, we read the safe raw tensors back into the model's state dict:

```python
from safetensors.torch import load_file

# Loading tensors without any execution risk
state_dict = load_file("model.safetensors")
model.load_state_dict(state_dict)

```

*(Note: If Needle uses custom NumPy arrays instead of standard PyTorch tensors, `safetensors` has a native `safetensors.numpy` module that works exactly the same way!)*

I'd be happy to open a PR to start making this transition if the maintainers are on board! Let me know what you think.

No project files or tests are identified. First locate the model serialization entry points that call torch.save, torch.load, or Python pickle, then determine whether the project uses PyTorch tensors or NumPy arrays. Done means the relevant model save/load paths use safetensors and the resulting behavior is verified by the project's tests.

Written by the indexing model from the issue text.