Update model.py
Browse files
model.py
CHANGED
|
@@ -2,11 +2,8 @@
|
|
| 2 |
# Enabling better customization for generalizable reward modeling
|
| 3 |
import torch
|
| 4 |
import torch.nn as nn
|
| 5 |
-
import os
|
| 6 |
from transformers import AutoModelForCausalLM
|
| 7 |
from trl import PreTrainedModelWrapper
|
| 8 |
-
from peft import PeftModel, PeftConfig
|
| 9 |
-
from safetensors import safe_open
|
| 10 |
|
| 11 |
|
| 12 |
class ValueHead(nn.Module):
|
|
|
|
| 2 |
# Enabling better customization for generalizable reward modeling
|
| 3 |
import torch
|
| 4 |
import torch.nn as nn
|
|
|
|
| 5 |
from transformers import AutoModelForCausalLM
|
| 6 |
from trl import PreTrainedModelWrapper
|
|
|
|
|
|
|
| 7 |
|
| 8 |
|
| 9 |
class ValueHead(nn.Module):
|