mirror of
https://github.com/wassname/Open-Assistant.git
synced 2026-09-09 11:15:08 +08:00
Merge branch 'LAION-AI:main' into main
This commit is contained in:
Alexander Goryunov
@@ -5,3 +5,6 @@
|
||||
*.egg-info
|
||||
__pycache__
|
||||
.DS_Store
|
||||
|
||||
# Generated files
|
||||
backend/oasst-openapi.json
|
||||
|
||||
@@ -179,3 +179,33 @@ if settings.DEBUG_USE_SEED_DATA:
|
||||
|
||||
|
||||
app.include_router(api_router, prefix=settings.API_V1_STR)
|
||||
|
||||
|
||||
def get_openapi_schema():
|
||||
return json.dumps(app.openapi())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Importing here so we don't import packages unnecessarily if we're
|
||||
# importing main as a module.
|
||||
import argparse
|
||||
import json
|
||||
|
||||
import uvicorn
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
|
||||
parser.add_argument(
|
||||
"--print-openapi-schema",
|
||||
help="Dumps the openapi schema to stdout",
|
||||
action=argparse.BooleanOptionalAction,
|
||||
)
|
||||
parser.add_argument("--host", help="The host to run the server")
|
||||
parser.add_argument("--port", help="The port to run the server")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.print_openapi_schema:
|
||||
print(get_openapi_schema())
|
||||
else:
|
||||
uvicorn.run(app, host=args.host, port=args.port)
|
||||
|
||||
@@ -153,7 +153,7 @@ def request_task(
|
||||
return task
|
||||
|
||||
|
||||
@router.post("/{task_id}/ack")
|
||||
@router.post("/{task_id}/ack", response_model=None)
|
||||
def tasks_acknowledge(
|
||||
*,
|
||||
db: Session = Depends(deps.get_db),
|
||||
@@ -179,10 +179,9 @@ def tasks_acknowledge(
|
||||
except Exception:
|
||||
logger.exception("Failed to acknowledge task.")
|
||||
raise OasstError("Failed to acknowledge task.", OasstErrorCode.TASK_ACK_FAILED)
|
||||
return {}
|
||||
|
||||
|
||||
@router.post("/{task_id}/nack")
|
||||
@router.post("/{task_id}/nack", response_model=None)
|
||||
def tasks_acknowledge_failure(
|
||||
*,
|
||||
db: Session = Depends(deps.get_db),
|
||||
@@ -204,7 +203,7 @@ def tasks_acknowledge_failure(
|
||||
raise OasstError("Failed to not acknowledge task.", OasstErrorCode.TASK_NACK_FAILED)
|
||||
|
||||
|
||||
@router.post("/interaction")
|
||||
@router.post("/interaction", response_model=protocol_schema.TaskDone)
|
||||
def tasks_interaction(
|
||||
*,
|
||||
db: Session = Depends(deps.get_db),
|
||||
|
||||
@@ -12,9 +12,17 @@ If you are unfamiliar with `hikari`, `lightbulb`, or `miru`, please refer to the
|
||||
|
||||
### Setup
|
||||
|
||||
To run the bot
|
||||
To run the bot:
|
||||
|
||||
Install dependency module `oasst-shared`
|
||||
|
||||
```bash
|
||||
cd oasst-shared
|
||||
pip install -e .
|
||||
```
|
||||
|
||||
```bash
|
||||
cd ../discord-bot
|
||||
cp .env.example .env
|
||||
|
||||
python -V # 3.10
|
||||
|
||||
@@ -53,14 +53,16 @@ class OasstApiClient:
|
||||
TaskType.done: protocol_schema.TaskDone,
|
||||
}
|
||||
|
||||
async def post(self, path: str, data: dict[str, t.Any]) -> dict[str, t.Any]:
|
||||
async def post(self, path: str, data: dict[str, t.Any]) -> Optional[dict[str, t.Any]]:
|
||||
"""Make a POST request to the backend."""
|
||||
logger.debug(f"POST {self.backend_url}{path} DATA: {data}")
|
||||
response = await self.session.post(f"{self.backend_url}{path}", json=data, headers={"X-API-Key": self.api_key})
|
||||
response.raise_for_status()
|
||||
return await response.json()
|
||||
|
||||
def _parse_task(self, data: dict[str, t.Any]) -> protocol_schema.Task:
|
||||
def _parse_task(self, data: Optional[dict[str, t.Any]]) -> protocol_schema.Task:
|
||||
if data is None:
|
||||
raise Exception("Cannot parse data as a task: data is none")
|
||||
task_type = TaskType(data.get("type"))
|
||||
|
||||
model = self.task_models_map.get(task_type)
|
||||
@@ -89,23 +91,22 @@ class OasstApiClient:
|
||||
logger.debug(f"Fetching random for user {user}")
|
||||
return await self.fetch_task(protocol_schema.TaskRequestType.random, user, collective)
|
||||
|
||||
async def ack_task(self, task_id: str | UUID, message_id: str):
|
||||
async def ack_task(self, task_id: str | UUID, message_id: str) -> None:
|
||||
"""Send an ACK for a task to the backend."""
|
||||
logger.debug(f"ACK task {task_id} with post {message_id}")
|
||||
req = protocol_schema.TaskAck(message_id=message_id)
|
||||
return await self.post(f"/api/v1/tasks/{task_id}/ack", data=req.dict())
|
||||
await self.post(f"/api/v1/tasks/{task_id}/ack", data=req.dict())
|
||||
|
||||
async def nack_task(self, task_id: str | UUID, reason: str):
|
||||
async def nack_task(self, task_id: str | UUID, reason: str) -> None:
|
||||
"""Send a NACK for a task to the backend."""
|
||||
logger.debug(f"NACK task {task_id} with reason {reason}")
|
||||
req = protocol_schema.TaskNAck(reason=reason)
|
||||
return await self.post(f"/api/v1/tasks/{task_id}/nack", data=req.dict())
|
||||
await self.post(f"/api/v1/tasks/{task_id}/nack", data=req.dict())
|
||||
|
||||
async def post_interaction(self, interaction: protocol_schema.Interaction) -> protocol_schema.Task:
|
||||
"""Send a completed task to the backend."""
|
||||
logger.debug(f"Interaction: {interaction}")
|
||||
resp = await self.post("/api/v1/tasks/interaction", data=interaction.dict())
|
||||
|
||||
return self._parse_task(resp)
|
||||
|
||||
async def close(self):
|
||||
|
||||
@@ -0,0 +1,2 @@
|
||||
pytest
|
||||
pytest-asyncio
|
||||
@@ -0,0 +1,52 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
from bot.api_client import OasstApiClient
|
||||
from oasst_shared.schemas import protocol as protocol_schema
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def oasst_api_client_mocked():
|
||||
client = OasstApiClient(backend_url="http://localhost:8080", api_key="123")
|
||||
yield client
|
||||
# TODO The fixture should close this connection, but there seems to be a bug
|
||||
# with async fixtures and pytest.
|
||||
# Since this only results in a warning, I'm leaving this for now.
|
||||
# await client.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("task_type", protocol_schema.TaskRequestType)
|
||||
async def test_can_fetch_task(task_type: protocol_schema.TaskRequestType, oasst_api_client_mocked: OasstApiClient):
|
||||
assert await oasst_api_client_mocked.fetch_task(task_type=task_type) is not None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_can_ack_task(oasst_api_client_mocked: OasstApiClient):
|
||||
await oasst_api_client_mocked.ack_task(task_id=uuid4(), message_id="123")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_can_nack_task(oasst_api_client_mocked: OasstApiClient):
|
||||
await oasst_api_client_mocked.nack_task(task_id=uuid4(), reason="bad task")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_can_post_interaction(oasst_api_client_mocked: OasstApiClient):
|
||||
assert (
|
||||
await oasst_api_client_mocked.post_interaction(
|
||||
protocol_schema.TextReplyToMessage(
|
||||
type="text_reply_to_message",
|
||||
message_id="123",
|
||||
user_message_id="321",
|
||||
text="This is my reply",
|
||||
user=protocol_schema.User(
|
||||
id="123",
|
||||
display_name="lomz",
|
||||
auth_method="discord",
|
||||
),
|
||||
)
|
||||
)
|
||||
is not None
|
||||
)
|
||||
@@ -27,6 +27,26 @@ services:
|
||||
timeout: 2s
|
||||
retries: 10
|
||||
|
||||
# Redis - caching + rate limiting on BE
|
||||
redis:
|
||||
image: redis
|
||||
restart: always
|
||||
ports:
|
||||
- 6379:6379
|
||||
healthcheck:
|
||||
test: ["CMD-SHELL", "redis-cli ping | grep PONG"]
|
||||
interval: 2s
|
||||
timeout: 2s
|
||||
retries: 10
|
||||
command: redis-server /usr/local/etc/redis/redis.conf
|
||||
volumes:
|
||||
- ./redis.conf:/usr/local/etc/redis/redis.conf
|
||||
# insights host - redis:6379
|
||||
redis-insights:
|
||||
image: redislabs/redisinsight:latest
|
||||
ports:
|
||||
- 8001:8001
|
||||
|
||||
# This DB is for Web Authentication and data caching.
|
||||
webdb:
|
||||
image: postgres
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
FROM python:3.10-slim-bullseye
|
||||
RUN mkdir /app
|
||||
COPY ./discord-bot/requirements.txt /requirements.txt
|
||||
RUN pip install -r requirements.txt
|
||||
WORKDIR /app
|
||||
COPY ./discord-bot /app
|
||||
CMD ["python", "bot.py"]
|
||||
COPY ./oasst-shared/oasst_shared /app/oasst_shared
|
||||
RUN pip install -r requirements.txt
|
||||
CMD ["python","-m","bot"]
|
||||
|
||||
@@ -0,0 +1,23 @@
|
||||
# Research
|
||||
|
||||
This page lists research papers that are relevant to the project.
|
||||
|
||||
## Automatically Generating Instruction Data for Training
|
||||
|
||||
This line of work is about significantly reducing the need for manually annotated data for the purpose of training [instruction-aligned](https://openai.com/blog/instruction-following/) language models.
|
||||
|
||||
### SELF-INSTRUCT: Aligning Language Model with Self Generated Instructions [[ArXiv](https://arxiv.org/pdf/2212.10560.pdf)], [[Github](https://github.com/yizhongw/self-instruct)].
|
||||
|
||||
> We introduce SELF-INSTRUCT, a framework for improving the instruction-following capabilities of pretrained language models by bootstrapping off its own generations.
|
||||
> Our pipeline generates instruction, input, and output samples from a language model, then prunes them before using them to finetune the original model.
|
||||
> Applying our method to vanilla GPT3, we demonstrate a 33% absolute improvement over the original model on SuperNaturalInstructions, on par with the performance of InstructGPT-0011, which is trained with private user data and human annotations.
|
||||
|
||||
### Tuning Language Models with (Almost) No Human Labor. [[ArXiv](https://arxiv.org/pdf/2212.09689.pdf)], [[Github](https://github.com/orhonovich/unnatural-instructions)].
|
||||
|
||||
> In this work, we introduce
|
||||
> Unnatural Instructions: a large dataset of creative and diverse instructions, collected with virtually no human labor.
|
||||
> We collect 64,000 examples by prompting a language model with three seed examples of instructions and eliciting a fourth.
|
||||
> This set is then expanded by prompting the model to rephrase each instruction, creating a total of approximately 240,000 examples of instructions, inputs, and outputs.
|
||||
> Experiments show that despite containing a fair amount of noise, training on Unnatural Instructions rivals the effectiveness of training
|
||||
> on open-source manually-curated datasets, surpassing the performance of models such as
|
||||
> T0++ and Tk-Instruct across various benchmarks.
|
||||
@@ -0,0 +1,53 @@
|
||||
# Sections to train Reward Model (RM)
|
||||
|
||||
Trainer code based on huggingface. Compatible with deepspeed or accelerate
|
||||
|
||||
Requirements
|
||||
|
||||
```
|
||||
wandb
|
||||
evaluate
|
||||
datasets
|
||||
transformers
|
||||
torch==1.12
|
||||
```
|
||||
|
||||
Start training reward model
|
||||
|
||||
```bash
|
||||
python trainer.py configs/electra-base-dis-webgpt.yml
|
||||
```
|
||||
|
||||
Additional axis labeling, this outputs a 4 summary quality evaluation metrics (score are normalized to 0-1 )
|
||||
|
||||
```bash
|
||||
python summary_quality_trainer.py configs/test-bloomz-560m-quality.yml
|
||||
```
|
||||
|
||||
The four summary are :
|
||||
|
||||
- overall
|
||||
|
||||
- accuracy
|
||||
|
||||
- coverage
|
||||
|
||||
- coherence
|
||||
|
||||
## Dataset
|
||||
|
||||
For now we only supports webgpt and summary dataset from OpenAI. Once open-asisstant dataset are available it will be added here.
|
||||
|
||||
## Model
|
||||
|
||||
Check out configs
|
||||
|
||||
```
|
||||
Open-Assistant/model/reward/instructor/configs/
|
||||
bloomz-560m.yml
|
||||
electra-base-dis-webgpt.yml
|
||||
galactica-125m.yml
|
||||
galactica-1b.yml
|
||||
```
|
||||
|
||||
You can add new huggingface model as you want.
|
||||
@@ -0,0 +1,19 @@
|
||||
Some other reward features we can use
|
||||
|
||||
0. Finish classifcation feature
|
||||
|
||||
1. Summaries from human feedback
|
||||
|
||||
- use `confidence` score into the RM learning, ensure the output rank score correlates with confidence
|
||||
|
||||
- each labeling has a labeling `note`, basically comments by labeler, not sure what else we can use
|
||||
|
||||
- ~~Use the score for "overall", "accuracy", "coverage", "coherence" from axis/evals to train an addition model (rank additional aspect of the policy model)~~
|
||||
|
||||
- this should be placed under experimental_dataset.py
|
||||
|
||||
2. Add support for anthropic dataset
|
||||
|
||||
- anthropic dataset is more like a conversation tree which is much complex than simply question-answer schema
|
||||
|
||||
- this is basically a MCTS from alphazero.
|
||||
@@ -0,0 +1,65 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
|
||||
classification based ranking
|
||||
|
||||
"""
|
||||
import json
|
||||
import os
|
||||
import random
|
||||
|
||||
from dataset import load_dataset
|
||||
from torch.utils.data import Dataset
|
||||
|
||||
from .utils import webgpt_return_format
|
||||
|
||||
|
||||
class WebGPTDataset(Dataset):
|
||||
def __init__(self, mode="train", index_cache="dataset/webgpt_train_idx.pt", additional_dataset=None) -> None:
|
||||
super().__init__()
|
||||
"""
|
||||
mode : train or val, used for validation purpose, has nothing to do with original split
|
||||
additional_dataset : a list of jsonline format with idx, question and texts (generate candidates)
|
||||
idx : must match the index you iterate from comparison enumerate order
|
||||
question : for validation purpose
|
||||
texts : list of K generate results from the question prompt
|
||||
"""
|
||||
os.makedirs("dataset", exist_ok=True)
|
||||
dataset = load_dataset("openai/webgpt_comparisons")
|
||||
self.dataset = []
|
||||
self.dataset_index = []
|
||||
for idx, row in enumerate(dataset["train"]):
|
||||
self.dataset.append(webgpt_return_format(row))
|
||||
|
||||
# since this dataset was generated from 176B GPT-3
|
||||
# we needed some more sample generated from the starting model
|
||||
# since this model must rank model generated by GPT-3 being better than your starting model
|
||||
self.sample_additional = False
|
||||
if additional_dataset is not None:
|
||||
self.sample_additional = True
|
||||
self.additional = {}
|
||||
with open(additional_dataset, "r") as f:
|
||||
for line in f:
|
||||
row = json.loads(line)
|
||||
if row["idx"] in self.dataset_index:
|
||||
self.additional[row["idx"]] = row["negatives"]
|
||||
if len(self.additional) != len(self.dataset_index):
|
||||
for match_idx in self.dataset_index:
|
||||
if match_idx in self.additional:
|
||||
continue
|
||||
|
||||
idx = match_idx - 900
|
||||
while idx not in self.additional:
|
||||
idx -= 1
|
||||
self.additional[match_idx] = self.additional[idx]
|
||||
|
||||
def __len__(self):
|
||||
return len(self.dataset)
|
||||
|
||||
def __getitem__(self, index):
|
||||
row = self.dataset[index]
|
||||
if not self.sample_additional:
|
||||
return row["question"], row["pos"], row["neg"]
|
||||
|
||||
gen_neg = random.choice(self.additional[self.dataset_index[index]])
|
||||
return row["question"], row["pos"], row["neg"], gen_neg
|
||||
@@ -0,0 +1,9 @@
|
||||
model_name: bigscience/bloomz-560m
|
||||
learning_rate: 3e-5
|
||||
gradient_accumulation_steps: 16
|
||||
per_device_train_batch_size: 2
|
||||
max_length: 600
|
||||
freeze_layer: 12
|
||||
num_train_epochs: 2
|
||||
datasets:
|
||||
- hfsummary
|
||||
@@ -0,0 +1,10 @@
|
||||
model_name: bigscience/bloomz-560m
|
||||
learning_rate: 3e-5
|
||||
gradient_accumulation_steps: 16
|
||||
per_device_train_batch_size: 2
|
||||
max_length: 600
|
||||
freeze_layer: 12
|
||||
num_train_epochs: 2
|
||||
datasets:
|
||||
- webgpt
|
||||
- hfsummary
|
||||
@@ -0,0 +1,3 @@
|
||||
model_name: google/electra-large-discriminator
|
||||
learning_rate: 3e-5
|
||||
max_length: 300
|
||||
@@ -0,0 +1,13 @@
|
||||
model_name: facebook/galactica-125m
|
||||
learning_rate: 1e-5
|
||||
gradient_checkpointing: false
|
||||
gradient_accumulation_steps: 32
|
||||
per_device_train_batch_size: 2
|
||||
warmup_steps: 600
|
||||
eval_steps: 200
|
||||
save_steps: 500
|
||||
max_length: 512
|
||||
num_train_epochs: 2
|
||||
datasets:
|
||||
- webgpt
|
||||
- hfsummary
|
||||
@@ -0,0 +1,14 @@
|
||||
model_name: facebook/galactica-1.3b
|
||||
learning_rate: 6e-6
|
||||
gradient_checkpointing: false
|
||||
gradient_accumulation_steps: 16
|
||||
per_device_train_batch_size: 2
|
||||
warmup_steps: 600
|
||||
freeze_layer: 20
|
||||
eval_steps: 200
|
||||
save_steps: 500
|
||||
max_length: 400
|
||||
num_train_epochs: 2
|
||||
datasets:
|
||||
- webgpt
|
||||
- hfsummary
|
||||
@@ -0,0 +1,14 @@
|
||||
model_name: facebook/galactica-125m
|
||||
learning_rate: 1e-5
|
||||
gradient_checkpointing: false
|
||||
gradient_accumulation_steps: 10
|
||||
per_device_train_batch_size: 6
|
||||
warmup_steps: 600
|
||||
loss: cls
|
||||
eval_steps: 200
|
||||
save_steps: 500
|
||||
max_length: 128
|
||||
num_train_epochs: 2
|
||||
datasets:
|
||||
- webgpt
|
||||
- hfsummary
|
||||
@@ -0,0 +1,100 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
HFSummary
|
||||
|
||||
I want to train a multi regression model on axis_evals dataset mainly we can estimate the score of these score
|
||||
|
||||
- {"overall": "6", "accuracy": "6", "coverage": "6", "coherence": "7"}
|
||||
|
||||
Should be better than just a preference score
|
||||
|
||||
"""
|
||||
from collections import defaultdict
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional, Union
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from datasets import load_dataset
|
||||
from torch.utils.data import Dataset
|
||||
from transformers.tokenization_utils_base import PaddingStrategy, PreTrainedTokenizerBase
|
||||
|
||||
|
||||
@dataclass
|
||||
class DataCollatorForSummaryScore:
|
||||
"""
|
||||
|
||||
Data collator that will dynamically pad the inputs for multiple choice received.
|
||||
|
||||
"""
|
||||
|
||||
tokenizer: PreTrainedTokenizerBase
|
||||
num_choices: int = 2
|
||||
padding: Union[bool, str, PaddingStrategy] = True
|
||||
max_length: Optional[int] = None
|
||||
pad_to_multiple_of: Optional[int] = None
|
||||
drop_token_type: bool = False # galactica
|
||||
|
||||
def __call__(self, batch):
|
||||
|
||||
features = []
|
||||
labels = []
|
||||
for feature, label in batch:
|
||||
features.append(feature)
|
||||
labels.append(label)
|
||||
|
||||
batch_feature = self.tokenizer.pad(
|
||||
features,
|
||||
padding=self.padding,
|
||||
max_length=self.max_length,
|
||||
pad_to_multiple_of=self.pad_to_multiple_of,
|
||||
return_tensors="pt",
|
||||
)
|
||||
if self.drop_token_type:
|
||||
batch_feature.pop("token_type_ids")
|
||||
# batch = {k: v.view(batch_size, self.num_choices, -1) for k, v in batch.items()}
|
||||
batch_feature["labels"] = torch.from_numpy(np.array(labels)).float()
|
||||
return batch_feature
|
||||
|
||||
|
||||
class HFSummaryQuality(Dataset):
|
||||
def __init__(self, split, tokenizer, max_length=300) -> None:
|
||||
super().__init__()
|
||||
assert split in ("validation", "test")
|
||||
dataset = load_dataset("Tristan/summarize_from_feedback", "axis")[split]
|
||||
self.max_length = max_length
|
||||
mean_scores = defaultdict(list)
|
||||
self.contexts = []
|
||||
self.responses = []
|
||||
self.labels = []
|
||||
for data in dataset:
|
||||
|
||||
if "article" in data["info"] and data["info"]["article"] is not None:
|
||||
context = data["info"]["article"]
|
||||
elif "post" in data["info"]:
|
||||
context = data["info"]["post"]
|
||||
self.contexts.append(context)
|
||||
|
||||
response = data["summary"]["text"]
|
||||
self.responses.append(response)
|
||||
self.labels.append(data["summary"]["axes"])
|
||||
for axis, score in data["summary"]["axes"].items():
|
||||
if score is not None:
|
||||
mean_scores[axis].append(score)
|
||||
|
||||
self.label2idx = {key: idx for idx, key in enumerate(mean_scores.keys())}
|
||||
self.label2mean = {key: np.mean(scores) for key, scores in mean_scores.items()}
|
||||
self.tokenizer = tokenizer
|
||||
print(self.label2idx)
|
||||
|
||||
def __len__(self):
|
||||
return len(self.responses)
|
||||
|
||||
def __getitem__(self, index):
|
||||
context = self.contexts[index]
|
||||
# return pairs of comparison
|
||||
response = self.responses[index]
|
||||
labels = np.zeros(len(self.label2idx))
|
||||
for key, score in self.labels[index].items():
|
||||
labels[self.label2idx[key]] = (self.label2mean[key] if score is None else score) / 10
|
||||
return self.tokenizer(context, response, truncation=True, max_length=self.max_length), labels
|
||||
@@ -0,0 +1,166 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
author: theblackcat102
|
||||
|
||||
Dataset output format from __getitem__
|
||||
|
||||
- question / prompt : string
|
||||
|
||||
- answers / rows : list of tuple pair. The first element in the tuple pair must be the positive pair (rank higher than the second element)
|
||||
|
||||
A list of rank based dataset for training using rank loss
|
||||
|
||||
Some nice features to have
|
||||
|
||||
[] support additional negative samples generated from other models.
|
||||
|
||||
For example we can use galactica-125m to generate a TLDR and assume it was
|
||||
inferior than the human perference one
|
||||
|
||||
|
||||
"""
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional, Union
|
||||
|
||||
import numpy as np
|
||||
from datasets import load_dataset
|
||||
from torch.utils.data import Dataset
|
||||
from transformers.tokenization_utils_base import PaddingStrategy, PreTrainedTokenizerBase
|
||||
|
||||
|
||||
@dataclass
|
||||
class DataCollatorForPairRank:
|
||||
"""
|
||||
|
||||
Data collator that will dynamically pad the inputs for multiple choice received.
|
||||
|
||||
"""
|
||||
|
||||
tokenizer: PreTrainedTokenizerBase
|
||||
num_choices: int = 2
|
||||
padding: Union[bool, str, PaddingStrategy] = True
|
||||
max_length: Optional[int] = None
|
||||
pad_to_multiple_of: Optional[int] = None
|
||||
drop_token_type: bool = False # galactica
|
||||
|
||||
def __call__(self, features):
|
||||
|
||||
flatten_features = []
|
||||
batch_size = 0
|
||||
for question, pairs in features:
|
||||
for (pos, neg) in pairs:
|
||||
flatten_features.append(self.tokenizer(question, pos, truncation=True, max_length=self.max_length))
|
||||
flatten_features.append(self.tokenizer(question, neg, truncation=True, max_length=self.max_length))
|
||||
batch_size += 1
|
||||
|
||||
batch = self.tokenizer.pad(
|
||||
flatten_features,
|
||||
padding=self.padding,
|
||||
max_length=self.max_length,
|
||||
pad_to_multiple_of=self.pad_to_multiple_of,
|
||||
return_tensors="pt",
|
||||
)
|
||||
if self.drop_token_type:
|
||||
batch.pop("token_type_ids")
|
||||
# batch = {k: v.view(batch_size, self.num_choices, -1) for k, v in batch.items()}
|
||||
return batch
|
||||
|
||||
|
||||
class WebGPT(Dataset):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
|
||||
dataset = load_dataset("openai/webgpt_comparisons")
|
||||
questions = {}
|
||||
# using prompt as our index will allows us
|
||||
# to add additional generated prompt later
|
||||
self.index2question = {}
|
||||
for row in dataset["train"]:
|
||||
question = row["question"]["full_text"]
|
||||
if question not in self.index2question:
|
||||
self.index2question[len(self.index2question)] = question
|
||||
|
||||
if question not in questions:
|
||||
questions[question] = []
|
||||
|
||||
if row["score_0"] > row["score_1"]:
|
||||
# not going to risk it
|
||||
questions[question].append((row["answer_0"], row["answer_1"]))
|
||||
else:
|
||||
questions[question].append((row["answer_1"], row["answer_0"]))
|
||||
|
||||
self.questions = questions
|
||||
|
||||
def __len__(self):
|
||||
return len(self.index2question)
|
||||
|
||||
def __getitem__(self, index):
|
||||
question = self.index2question[index]
|
||||
rows = self.questions[question]
|
||||
# optimize the format later
|
||||
return question, rows
|
||||
|
||||
|
||||
class HFSummary(Dataset):
|
||||
"""
|
||||
Human feedback data from OpenAI
|
||||
https://github.com/openai/summarize-from-feedback
|
||||
|
||||
labeling method : pair comparison, 0 or 1
|
||||
|
||||
"""
|
||||
|
||||
def __init__(self, split="train", conf_threshold=-1, max_comparison_per_sample=3) -> None:
|
||||
super().__init__()
|
||||
assert split in ("train", "valid1", "valid2", "test")
|
||||
summaries = {}
|
||||
# using prompt as our index will allows us
|
||||
# to add additional generated prompt later
|
||||
self.index2summary = {}
|
||||
self.max_comparison_per_sample = max_comparison_per_sample
|
||||
major_split = split if "train" == split else "validation"
|
||||
dataset = load_dataset("Tristan/summarize_from_feedback", "comparisons")[major_split]
|
||||
for data in dataset:
|
||||
if (
|
||||
"extra" in data
|
||||
and "confidence" in data["extra"]
|
||||
and data["extra"]["confidence"] is not None
|
||||
and conf_threshold > data["extra"]["confidence"]
|
||||
):
|
||||
print("skipping {}".format(data["info"]["id"]))
|
||||
continue
|
||||
|
||||
if split != "train" and split != data["split"]:
|
||||
continue
|
||||
|
||||
if "article" in data["info"] and data["info"]["article"] is not None:
|
||||
context = data["info"]["article"]
|
||||
elif "post" in data["info"]:
|
||||
context = data["info"]["post"]
|
||||
|
||||
if context not in self.index2summary:
|
||||
self.index2summary[len(self.index2summary)] = context
|
||||
|
||||
if context not in summaries:
|
||||
summaries[context] = []
|
||||
|
||||
pos, neg = (0, 1) if data["choice"] == 0 else (1, 0)
|
||||
summaries[context].append((data["summaries"][pos]["text"], data["summaries"][neg]["text"]))
|
||||
|
||||
self.summaries = summaries
|
||||
|
||||
self.postfix_prompt = " TLDR;"
|
||||
|
||||
def __len__(self):
|
||||
return len(self.index2summary)
|
||||
|
||||
def __getitem__(self, index):
|
||||
context = self.index2summary[index]
|
||||
# return pairs of comparison
|
||||
rows = self.summaries[context]
|
||||
# pair very big
|
||||
# we are going to do some sampling
|
||||
# not optimal but good for now
|
||||
valid_idx = np.random.choice(len(rows), self.max_comparison_per_sample)
|
||||
# optimize the format later
|
||||
return context + self.postfix_prompt, [r for idx, r in enumerate(rows) if idx in valid_idx]
|
||||
@@ -0,0 +1,6 @@
|
||||
datasets==2.8.0
|
||||
evaluate==0.4.0
|
||||
scikit-learn==1.2.0
|
||||
torch==1.12.1+cu116
|
||||
transformers==4.25.1
|
||||
wandb==0.13.7
|
||||
@@ -0,0 +1,156 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
import os
|
||||
from argparse import ArgumentParser
|
||||
from typing import Any, Callable, Dict, List, Optional, Tuple, Union
|
||||
|
||||
import evaluate
|
||||
import numpy as np
|
||||
import torch
|
||||
from experimental_dataset import DataCollatorForSummaryScore, HFSummaryQuality
|
||||
from torch import nn
|
||||
from torch.utils.data import Dataset
|
||||
from transformers import (
|
||||
AutoModelForSequenceClassification,
|
||||
DataCollator,
|
||||
EvalPrediction,
|
||||
PreTrainedModel,
|
||||
PreTrainedTokenizerBase,
|
||||
Trainer,
|
||||
TrainerCallback,
|
||||
TrainingArguments,
|
||||
)
|
||||
from utils import argument_parsing, freeze_top_n_layers, get_tokenizer
|
||||
|
||||
os.environ["WANDB_PROJECT"] = "quality-scoring"
|
||||
|
||||
parser = ArgumentParser()
|
||||
parser.add_argument("config", type=str)
|
||||
|
||||
accuracy = evaluate.load("mse")
|
||||
|
||||
|
||||
def compute_metrics(eval_pred):
|
||||
predictions, labels = eval_pred
|
||||
return accuracy.compute(predictions=predictions.flatten(), references=labels.flatten())
|
||||
|
||||
|
||||
class QualityTrainer(Trainer):
|
||||
def __init__(
|
||||
self,
|
||||
model: Union[PreTrainedModel, nn.Module] = None,
|
||||
args: TrainingArguments = None,
|
||||
data_collator: Optional[DataCollator] = None,
|
||||
train_dataset: Optional[Dataset] = None,
|
||||
eval_dataset: Optional[Dataset] = None,
|
||||
tokenizer: Optional[PreTrainedTokenizerBase] = None,
|
||||
model_init: Callable[[], PreTrainedModel] = None,
|
||||
compute_metrics: Optional[Callable[[EvalPrediction], Dict]] = None,
|
||||
callbacks: Optional[List[TrainerCallback]] = None,
|
||||
optimizers: Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None),
|
||||
preprocess_logits_for_metrics: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] = None,
|
||||
):
|
||||
super().__init__(
|
||||
model,
|
||||
args,
|
||||
data_collator,
|
||||
train_dataset,
|
||||
eval_dataset,
|
||||
tokenizer,
|
||||
model_init,
|
||||
compute_metrics,
|
||||
callbacks,
|
||||
optimizers,
|
||||
preprocess_logits_for_metrics,
|
||||
)
|
||||
self.loss_fct = nn.L1Loss()
|
||||
self.sigmoid = nn.Sigmoid()
|
||||
|
||||
def compute_loss(self, model, inputs, return_outputs=False):
|
||||
labels = inputs.pop("labels")
|
||||
# forward pass
|
||||
outputs = model(**inputs)
|
||||
logits = self.sigmoid(outputs.get("logits"))
|
||||
loss = self.loss_fct(logits, labels)
|
||||
|
||||
return (loss, outputs) if return_outputs else loss
|
||||
|
||||
def _compute_loss(self, model, inputs):
|
||||
inputs = self._prepare_inputs(inputs)
|
||||
labels = inputs.pop("labels")
|
||||
outputs = model(**inputs)
|
||||
logits = self.sigmoid(outputs.get("logits"))
|
||||
loss = self.loss_fct(logits, labels)
|
||||
|
||||
return loss, logits
|
||||
|
||||
def prediction_step(
|
||||
self,
|
||||
model: nn.Module,
|
||||
inputs: Dict[str, Union[torch.Tensor, Any]],
|
||||
prediction_loss_only: bool,
|
||||
ignore_keys: Optional[List[str]] = None,
|
||||
) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]:
|
||||
|
||||
with torch.no_grad():
|
||||
# compute loss on predict data
|
||||
loss, logits = self._compute_loss(model, inputs)
|
||||
|
||||
loss = loss.mean().detach()
|
||||
labels = inputs["labels"]
|
||||
if self.args.prediction_loss_only:
|
||||
return (loss, None, None)
|
||||
|
||||
return (loss, logits, labels)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
training_conf = argument_parsing(parser)
|
||||
|
||||
model_name = training_conf["model_name"]
|
||||
tokenizer = get_tokenizer(model_name)
|
||||
collate_fn = DataCollatorForSummaryScore(
|
||||
tokenizer, max_length=training_conf["max_length"], drop_token_type="galactica" in model_name
|
||||
)
|
||||
train = HFSummaryQuality(split="validation", tokenizer=tokenizer, max_length=training_conf["max_length"])
|
||||
eval = HFSummaryQuality(split="test", tokenizer=tokenizer, max_length=training_conf["max_length"])
|
||||
model = AutoModelForSequenceClassification.from_pretrained(
|
||||
model_name, num_labels=len(train.label2idx), problem_type="regression"
|
||||
)
|
||||
|
||||
if "freeze_layer" in training_conf:
|
||||
num_layer = training_conf["freeze_layer"]
|
||||
model = freeze_top_n_layers(model, num_layer)
|
||||
model_parameters = filter(lambda p: p.requires_grad, model.parameters())
|
||||
params = sum([np.prod(p.size()) for p in model_parameters])
|
||||
print("Number of trainable : {}M".format(int(params / 1e6)))
|
||||
|
||||
args = TrainingArguments(
|
||||
output_dir=f"{model_name}-finetuned",
|
||||
num_train_epochs=training_conf["num_train_epochs"],
|
||||
warmup_steps=500,
|
||||
learning_rate=training_conf["learning_rate"],
|
||||
# half_precision_backend="apex",
|
||||
fp16=True,
|
||||
gradient_checkpointing=training_conf["gradient_checkpointing"],
|
||||
gradient_accumulation_steps=training_conf["gradient_accumulation_steps"],
|
||||
per_device_train_batch_size=training_conf["per_device_train_batch_size"],
|
||||
per_device_eval_batch_size=training_conf["per_device_eval_batch_size"],
|
||||
weight_decay=0.01,
|
||||
max_grad_norm=2.0,
|
||||
logging_steps=10,
|
||||
save_total_limit=4,
|
||||
evaluation_strategy="steps",
|
||||
eval_steps=training_conf["eval_steps"],
|
||||
save_steps=1000,
|
||||
report_to="wandb",
|
||||
)
|
||||
trainer = QualityTrainer(
|
||||
model,
|
||||
args,
|
||||
train_dataset=train,
|
||||
eval_dataset=eval,
|
||||
data_collator=collate_fn,
|
||||
tokenizer=tokenizer,
|
||||
compute_metrics=compute_metrics,
|
||||
)
|
||||
trainer.train()
|
||||
@@ -0,0 +1,41 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
from experimental_dataset import DataCollatorForSummaryScore, HFSummaryQuality
|
||||
from rank_datasets import DataCollatorForPairRank, HFSummary, WebGPT
|
||||
from torch.utils.data import DataLoader
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
|
||||
def test_hfsummary():
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained("bigscience/mt0-large")
|
||||
collate_fn = DataCollatorForPairRank(tokenizer, max_length=200)
|
||||
dataset = HFSummary("train")
|
||||
print(len(dataset))
|
||||
dataloader = DataLoader(dataset, collate_fn=collate_fn, batch_size=8)
|
||||
for batch in dataloader:
|
||||
batch["input_ids"].shape
|
||||
|
||||
|
||||
def test_webgpt():
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained("bigscience/mt0-large")
|
||||
collate_fn = DataCollatorForPairRank(tokenizer, max_length=200)
|
||||
dataset = WebGPT()
|
||||
dataloader = DataLoader(dataset, collate_fn=collate_fn, batch_size=32)
|
||||
for batch in dataloader:
|
||||
print(batch["input_ids"].shape)
|
||||
|
||||
|
||||
def test_hf_quality():
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained("bigscience/mt0-large")
|
||||
collate_fn = DataCollatorForSummaryScore(tokenizer, max_length=200)
|
||||
dataset = HFSummaryQuality("validation", tokenizer)
|
||||
dataloader = DataLoader(dataset, collate_fn=collate_fn, batch_size=32)
|
||||
for batch in dataloader:
|
||||
print(batch["input_ids"].shape)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_hf_quality()
|
||||
# test_webgpt()
|
||||
@@ -0,0 +1,186 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
import os
|
||||
from argparse import ArgumentParser
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Callable, Dict, List, Optional, Tuple, Union
|
||||
|
||||
import evaluate
|
||||
import numpy as np
|
||||
import torch
|
||||
from rank_datasets import DataCollatorForPairRank, HFSummary, WebGPT
|
||||
from torch import nn
|
||||
from torch.utils.data import ConcatDataset, Dataset
|
||||
from transformers import (
|
||||
AutoModelForSequenceClassification,
|
||||
DataCollator,
|
||||
EvalPrediction,
|
||||
PreTrainedModel,
|
||||
PreTrainedTokenizerBase,
|
||||
Trainer,
|
||||
TrainerCallback,
|
||||
TrainingArguments,
|
||||
)
|
||||
from utils import argument_parsing, freeze_top_n_layers, get_tokenizer, train_val_dataset
|
||||
|
||||
os.environ["WANDB_PROJECT"] = "reward-model"
|
||||
|
||||
accuracy = evaluate.load("accuracy")
|
||||
parser = ArgumentParser()
|
||||
parser.add_argument("config", type=str)
|
||||
|
||||
|
||||
@dataclass
|
||||
class CustomTrainingArguments(TrainingArguments):
|
||||
loss_function: str = "rank"
|
||||
|
||||
|
||||
def compute_metrics(eval_pred):
|
||||
predictions, _ = eval_pred
|
||||
predictions = np.argmax(predictions, axis=1)
|
||||
return accuracy.compute(predictions=predictions, references=[0] * predictions.shape[0])
|
||||
|
||||
|
||||
class RankLoss(nn.Module):
|
||||
def __init__(self, eps=1e-8) -> None:
|
||||
super().__init__()
|
||||
self.eps = eps
|
||||
self.log_sigmoid = nn.LogSigmoid()
|
||||
|
||||
def forward(self, pos, neg):
|
||||
return -self.log_sigmoid(pos - neg + self.eps).mean()
|
||||
|
||||
|
||||
class RankTrainer(Trainer):
|
||||
def __init__(
|
||||
self,
|
||||
model: Union[PreTrainedModel, nn.Module] = None,
|
||||
args: TrainingArguments = None,
|
||||
data_collator: Optional[DataCollator] = None,
|
||||
train_dataset: Optional[Dataset] = None,
|
||||
eval_dataset: Optional[Dataset] = None,
|
||||
tokenizer: Optional[PreTrainedTokenizerBase] = None,
|
||||
model_init: Callable[[], PreTrainedModel] = None,
|
||||
compute_metrics: Optional[Callable[[EvalPrediction], Dict]] = None,
|
||||
callbacks: Optional[List[TrainerCallback]] = None,
|
||||
optimizers: Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None),
|
||||
preprocess_logits_for_metrics: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] = None,
|
||||
):
|
||||
super().__init__(
|
||||
model,
|
||||
args,
|
||||
data_collator,
|
||||
train_dataset,
|
||||
eval_dataset,
|
||||
tokenizer,
|
||||
model_init,
|
||||
compute_metrics,
|
||||
callbacks,
|
||||
optimizers,
|
||||
preprocess_logits_for_metrics,
|
||||
)
|
||||
self.loss_fct = RankLoss() if args.loss_function == "rank" else nn.CrossEntropyLoss()
|
||||
self.loss_function = args.loss_function
|
||||
|
||||
def compute_loss(self, model, inputs, return_outputs=False):
|
||||
# forward pass
|
||||
outputs = model(**inputs)
|
||||
logits = outputs.get("logits").view(-1, 2)
|
||||
if self.loss_function == "rank":
|
||||
loss = self.loss_fct(logits[:, 0], logits[:, 1])
|
||||
else:
|
||||
loss = self.loss_fct(logits, torch.zeros(logits.shape[0], device=logits.device, dtype=torch.long))
|
||||
|
||||
return (loss, outputs) if return_outputs else loss
|
||||
|
||||
def _compute_loss(self, model, inputs):
|
||||
inputs = self._prepare_inputs(inputs)
|
||||
outputs = model(**inputs)
|
||||
logits = outputs.get("logits").view(-1, 2)
|
||||
if self.loss_function == "rank":
|
||||
loss = self.loss_fct(logits[:, 0], logits[:, 1])
|
||||
else:
|
||||
loss = self.loss_fct(logits, torch.zeros(logits.shape[0], device=logits.device, dtype=torch.long))
|
||||
|
||||
return loss, logits
|
||||
|
||||
def prediction_step(
|
||||
self,
|
||||
model: nn.Module,
|
||||
inputs: Dict[str, Union[torch.Tensor, Any]],
|
||||
prediction_loss_only: bool,
|
||||
ignore_keys: Optional[List[str]] = None,
|
||||
) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]:
|
||||
|
||||
with torch.no_grad():
|
||||
# compute loss on predict data
|
||||
loss, logits = self._compute_loss(model, inputs)
|
||||
|
||||
loss = loss.mean().detach()
|
||||
labels = torch.zeros(logits.shape[0], device=logits.device, dtype=torch.long)
|
||||
if self.args.prediction_loss_only:
|
||||
return (loss, None, None)
|
||||
|
||||
return (loss, logits, labels)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
training_conf = argument_parsing(parser)
|
||||
|
||||
model_name = training_conf["model_name"]
|
||||
model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=1, problem_type="regression")
|
||||
if "freeze_layer" in training_conf:
|
||||
num_layer = training_conf["freeze_layer"]
|
||||
model = freeze_top_n_layers(model, num_layer)
|
||||
model_parameters = filter(lambda p: p.requires_grad, model.parameters())
|
||||
params = sum([np.prod(p.size()) for p in model_parameters])
|
||||
print("Number of trainable : {}M".format(int(params / 1e6)))
|
||||
|
||||
tokenizer = get_tokenizer(model_name)
|
||||
args = CustomTrainingArguments(
|
||||
output_dir=f"{model_name}-finetuned",
|
||||
num_train_epochs=training_conf["num_train_epochs"],
|
||||
warmup_steps=500,
|
||||
loss_function=training_conf["loss"],
|
||||
learning_rate=training_conf["learning_rate"],
|
||||
# half_precision_backend="apex",
|
||||
fp16=True,
|
||||
gradient_checkpointing=training_conf["gradient_checkpointing"],
|
||||
gradient_accumulation_steps=training_conf["gradient_accumulation_steps"],
|
||||
per_device_train_batch_size=training_conf["per_device_train_batch_size"],
|
||||
per_device_eval_batch_size=training_conf["per_device_eval_batch_size"],
|
||||
weight_decay=0.01,
|
||||
max_grad_norm=2.0,
|
||||
logging_steps=10,
|
||||
save_total_limit=4,
|
||||
evaluation_strategy="steps",
|
||||
eval_steps=training_conf["eval_steps"],
|
||||
save_steps=1000,
|
||||
report_to="wandb",
|
||||
)
|
||||
train_datasets, evals = [], {}
|
||||
if "webgpt" in training_conf["datasets"]:
|
||||
web_dataset = WebGPT()
|
||||
train, eval = train_val_dataset(web_dataset)
|
||||
train_datasets.append(train)
|
||||
evals["webgpt"] = eval
|
||||
if "hfsummary" in training_conf["datasets"]:
|
||||
sum_train = HFSummary(split="train")
|
||||
train_datasets.append(sum_train)
|
||||
sum_eval = HFSummary(split="valid1")
|
||||
assert len(sum_eval) > 0
|
||||
evals["hfsummary"] = sum_eval
|
||||
train = ConcatDataset(train_datasets)
|
||||
collate_fn = DataCollatorForPairRank(
|
||||
tokenizer, max_length=training_conf["max_length"], drop_token_type="galactica" in model_name
|
||||
)
|
||||
assert len(evals) > 0
|
||||
trainer = RankTrainer(
|
||||
model,
|
||||
args,
|
||||
train_dataset=train,
|
||||
eval_dataset=eval,
|
||||
data_collator=collate_fn,
|
||||
tokenizer=tokenizer,
|
||||
compute_metrics=compute_metrics,
|
||||
)
|
||||
trainer.train()
|
||||
@@ -0,0 +1,100 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
import re
|
||||
|
||||
import yaml
|
||||
from sklearn.model_selection import train_test_split
|
||||
from torch.utils.data import Subset
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
re_reference_remove = re.compile(r"\[([0-9])+\]|\[([0-9])+,([0-9])+\]")
|
||||
|
||||
|
||||
def webgpt_return_format(row):
|
||||
if row["score_0"] >= row["score_1"]:
|
||||
# remove this to prevent information leak, since we are not using reference
|
||||
return {
|
||||
"question": row["question"]["full_text"],
|
||||
"pos": re_reference_remove.sub("", row["answer_0"]),
|
||||
"neg": re_reference_remove.sub("", row["answer_1"]),
|
||||
}
|
||||
|
||||
return {
|
||||
"question": row["question"]["full_text"],
|
||||
"pos": re_reference_remove.sub("", row["answer_1"]),
|
||||
"neg": re_reference_remove.sub("", row["answer_0"]),
|
||||
}
|
||||
|
||||
|
||||
def get_tokenizer(tokenizer_name):
|
||||
tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
|
||||
if "galactica" in tokenizer_name:
|
||||
tokenizer.add_special_tokens({"pad_token": "<pad>", "eos_token": "</s>"})
|
||||
|
||||
return tokenizer
|
||||
|
||||
|
||||
def train_val_dataset(dataset, val_split=0.2):
|
||||
train_idx, val_idx = train_test_split(
|
||||
list(range(len(dataset))), test_size=val_split, random_state=666, shuffle=True
|
||||
)
|
||||
# [3879, 11479, 8341, 9177, 10798, 18177, 5735, 15669, 4837, 2760]
|
||||
print(val_idx[:10])
|
||||
# [13582, 5919, 11875, 7373, 19135, 13706, 8555, 15788, 15005, 15209]
|
||||
print(train_idx[:10])
|
||||
return Subset(dataset, train_idx), Subset(dataset, val_idx)
|
||||
|
||||
|
||||
def freeze_top_n_layers(model, target_layers):
|
||||
# its possible we can simply detect which module is a ModuleList
|
||||
# and simply freeze the module without doing string parsing
|
||||
for name, param in model.named_parameters():
|
||||
if "embed" in name:
|
||||
param.requires_grad = False
|
||||
elif ".layer" in name or ".h." in name:
|
||||
tokens = name.split(".")
|
||||
idx = 0
|
||||
for token in tokens:
|
||||
if "layer" in token or token == "h":
|
||||
break
|
||||
idx += 1
|
||||
if idx >= len(tokens):
|
||||
continue
|
||||
|
||||
layer_ = int(tokens[idx + 1])
|
||||
if layer_ < target_layers:
|
||||
# print('freeze ', layer_, name)
|
||||
param.requires_grad = False
|
||||
return model
|
||||
|
||||
|
||||
def argument_parsing(parser):
|
||||
default_params = {
|
||||
"num_train_epochs": 4,
|
||||
"learning_rate": 3e-5,
|
||||
"eval_steps": 500,
|
||||
"loss": "rank",
|
||||
"max_length": 440,
|
||||
"per_device_eval_batch_size": 5,
|
||||
"per_device_train_batch_size": 8,
|
||||
"gradient_accumulation_steps": 8,
|
||||
"gradient_checkpointing": False,
|
||||
"datasets": ["webgpt"],
|
||||
}
|
||||
args = parser.parse_args()
|
||||
with open(args.config, "r", encoding="utf-8") as f:
|
||||
training_conf = yaml.safe_load(f.read())
|
||||
|
||||
params = {**default_params, **training_conf}
|
||||
params["gradient_accumulation_steps"] = int(params["gradient_accumulation_steps"])
|
||||
params["num_train_epochs"] = int(params["num_train_epochs"])
|
||||
params["per_device_train_batch_size"] = int(params["per_device_train_batch_size"])
|
||||
params["learning_rate"] = float(params["learning_rate"])
|
||||
return params
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
from transformers import AutoModelForSequenceClassification
|
||||
|
||||
model = AutoModelForSequenceClassification.from_pretrained("bigscience/bloomz-560m")
|
||||
freeze_top_n_layers(model, 10)
|
||||
print(model.state_dict().keys())
|
||||
@@ -5,7 +5,7 @@ from typing import Literal, Optional, Union
|
||||
from uuid import UUID, uuid4
|
||||
|
||||
import pydantic
|
||||
from pydantic import BaseModel
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
class TaskRequestType(str, enum.Enum):
|
||||
@@ -56,7 +56,9 @@ class TaskRequest(BaseModel):
|
||||
"""The frontend asks the backend for a task."""
|
||||
|
||||
type: TaskRequestType = TaskRequestType.random
|
||||
user: Optional[User] = None
|
||||
# Must use Field(..., nullable=True) to indicate to the OpenAPI schema that
|
||||
# this is optional. https://github.com/pydantic/pydantic/issues/1270
|
||||
user: Optional[User] = Field(None, nullable=True)
|
||||
collective: bool = False
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,2 @@
|
||||
maxmemory 100mb
|
||||
maxmemory-policy allkeys-lru
|
||||
+41
@@ -0,0 +1,41 @@
|
||||
#!/usr/bin/env bash
|
||||
parent_path=$( cd "$(dirname "${BASH_SOURCE[0]}")" ; pwd -P )
|
||||
|
||||
# switch to backend directory
|
||||
pushd "$parent_path/../../backend"
|
||||
|
||||
MOCK_SERVER_PORT=8080
|
||||
OPENAPI_JSON_FILE_NAME=openapi.json
|
||||
|
||||
echo "Generating OpenAPI schema..."
|
||||
python -m main --print-openapi-schema > $OPENAPI_JSON_FILE_NAME
|
||||
echo "Done!"
|
||||
|
||||
# If oasst-mock-backend docker container is already running,
|
||||
# just restart it
|
||||
if [ "$(docker ps -q -f name=oasst-mock-backend)" ]; then
|
||||
echo "oasst-mock-backend container exists, restarting..."
|
||||
docker restart oasst-mock-backend
|
||||
else
|
||||
echo "Creating new oasst-mock-backend container..."
|
||||
docker run --init --rm -d \
|
||||
--name oasst-mock-backend \
|
||||
-p $MOCK_SERVER_PORT:4010 \
|
||||
-v $(pwd):/tmp \
|
||||
-P stoplight/prism:4 \
|
||||
mock -h 0.0.0.0 "/tmp/$OPENAPI_JSON_FILE_NAME"
|
||||
fi
|
||||
|
||||
echo "Waiting for server to be live..."
|
||||
curl --retry-all-errors --retry 5 localhost:$MOCK_SERVER_PORT
|
||||
echo ""
|
||||
|
||||
# if return code is successful, print successful response
|
||||
if [ $? -eq 0 ]; then
|
||||
echo "Mock server is running at localhost:$MOCK_SERVER_PORT"
|
||||
else
|
||||
echo "Mock server failed to start"
|
||||
fi
|
||||
|
||||
|
||||
popd
|
||||
Executable
+9
@@ -0,0 +1,9 @@
|
||||
#!/usr/bin/env bash
|
||||
parent_path=$( cd "$(dirname "${BASH_SOURCE[0]}")" ; pwd -P )
|
||||
|
||||
# switch to backend directory
|
||||
pushd "$parent_path/../../discord-bot"
|
||||
|
||||
pytest .
|
||||
|
||||
popd
|
||||
Reference in New Issue
Block a user