01bb07dba8
FocalLoRA, Should-It-Be-Executed-Or-Processed, and topicattack were nested git repos (with an inner FocalLoRA/data/FocalLoRA/.git as well). Drop their .git history and track the contents directly in this repo instead of as submodules/gitlinks. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
124 lines
3.8 KiB
Python
124 lines
3.8 KiB
Python
|
|
import os
|
|
import json
|
|
import io
|
|
import random
|
|
import sys
|
|
|
|
import numpy as np
|
|
import torch.nn.functional as F
|
|
from transformers import AutoTokenizer, AutoModel
|
|
import torch
|
|
|
|
class Logger(object):
|
|
|
|
def __init__(self, log_path, on=True):
|
|
self.log_path = log_path
|
|
self.on = on
|
|
|
|
if self.on:
|
|
while os.path.isfile(self.log_path):
|
|
self.log_path += '+'
|
|
|
|
def log(self, string, newline=True, force=False):
|
|
if self.on or force:
|
|
with open(self.log_path, 'a') as logf:
|
|
string = str(string)
|
|
logf.write(string)
|
|
if newline: logf.write('\n')
|
|
|
|
sys.stdout.write(string)
|
|
if newline: sys.stdout.write('\n')
|
|
sys.stdout.flush()
|
|
|
|
def log_probs_from_logits(logits: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
|
|
log_probs = F.log_softmax(logits, dim=-1)
|
|
log_probs_labels = log_probs.gather(dim=-1, index=labels.unsqueeze(-1))
|
|
return log_probs_labels.squeeze(-1)
|
|
|
|
|
|
def get_sp_tokens(args):
|
|
sp_tokens = dict()
|
|
for key in ("bos_token", "eos_token", "pad_token", "unk_token"):
|
|
sp_token = getattr(args, key, None)
|
|
if sp_token is not None:
|
|
sp_tokens[key] = sp_token
|
|
return sp_tokens
|
|
|
|
def get_tokenizer(pretrain, model, padding_side="left", args=None, use_fast=True):
|
|
sp_tokens = get_sp_tokens(args)
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained(pretrain, trust_remote_code=True, **sp_tokens)
|
|
tokenizer.padding_side = padding_side
|
|
# NOTE: When enable vLLM, do not resize_token_embeddings, or the vocab size will mismatch with vLLM.
|
|
# https://github.com/facebookresearch/llama-recipes/pull/196
|
|
|
|
|
|
if "mistral" in pretrain.lower():
|
|
template_tokenizer = AutoTokenizer.from_pretrained('HuggingFaceH4/zephyr-7b-beta', trust_remote_code=True)
|
|
tokenizer.apply_chat_template = template_tokenizer.apply_chat_template
|
|
|
|
elif "llama" in pretrain.lower():
|
|
tempalte_tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct", trust_remote_code=True)
|
|
tokenizer.apply_chat_template = tempalte_tokenizer.apply_chat_template
|
|
# tokenizer.eos_token_id = 128001
|
|
# tokenizer.eos_token = '<|end_of_text|>'
|
|
|
|
if tokenizer.pad_token is None:
|
|
tokenizer.pad_token = tokenizer.eos_token
|
|
tokenizer.pad_token_id = tokenizer.eos_token_id
|
|
model.config.pad_token_id = tokenizer.pad_token_id
|
|
|
|
return tokenizer
|
|
|
|
def set_seeds(args):
|
|
random.seed(args.seed)
|
|
np.random.seed(args.seed)
|
|
torch.manual_seed(args.seed)
|
|
|
|
def _make_w_io_base(f, mode: str):
|
|
if not isinstance(f, io.IOBase):
|
|
f_dirname = os.path.dirname(f)
|
|
if f_dirname != "":
|
|
os.makedirs(f_dirname, exist_ok=True)
|
|
f = open(f, mode=mode)
|
|
return f
|
|
|
|
|
|
def _make_r_io_base(f, mode: str):
|
|
if not isinstance(f, io.IOBase):
|
|
f = open(f, mode=mode)
|
|
return f
|
|
|
|
|
|
def jdump(obj, f, mode="w", indent=4, default=str):
|
|
"""Dump a str or dictionary to a file in json format.
|
|
|
|
Args:
|
|
obj: An object to be written.
|
|
f: A string path to the location on disk.
|
|
mode: Mode for opening the file.
|
|
indent: Indent for storing json dictionaries.
|
|
default: A function to handle non-serializable entries; defaults to `str`.
|
|
"""
|
|
f = _make_w_io_base(f, mode)
|
|
if isinstance(obj, (dict, list)):
|
|
json.dump(obj, f, indent=indent, default=default)
|
|
elif isinstance(obj, str):
|
|
f.write(obj)
|
|
else:
|
|
raise ValueError(f"Unexpected type: {type(obj)}")
|
|
f.close()
|
|
|
|
|
|
def jload(f, mode="r"):
|
|
"""Load a .json file into a dictionary."""
|
|
f = _make_r_io_base(f, mode)
|
|
jdict = json.load(f)
|
|
f.close()
|
|
return jdict
|
|
|
|
|
|
def load_text(path):
|
|
with open(path, "r") as f:
|
|
return f.read() |