RAM_Altman_GPT / src /config.py
brunuu's picture
1.0
7de15fc verified
Raw History Blame Contribute Delete
857 Bytes
from pathlib import Path
import torch
import random
PROJECT_ROOT = Path(__file__).resolve().parent.parent
DATA_FILE = str(PROJECT_ROOT / "data" / "ram_altman.txt")
OUTPUT_DIR = str(PROJECT_ROOT / "RAM_Altman_GPT")
END_TOKEN = "<|i_have_enough_ram|>"
TARGET_PARAMS = 100_000
EMBEDDING_SIZE = 40
NUM_LAYERS = 3
NUM_HEADS = 4
VOCAB_SIZE = 640
STEPS = 30000
BATCH_SIZE = 32
LEARNING_RATE = 4e-3
CHAT_TEMPLATE = (
"{% for message in messages %}"
"{% if message['role'] == 'user' %}User: {{ message['content'] }}\n"
"{% elif message['role'] == 'assistant' %}RAM-ALTMAN: {{ message['content'] }}<|i_have_enough_ram|>\n"
"{% endif %}"
"{% endfor %}"
"{% if add_generation_prompt %}RAM-ALTMAN:{% endif %}"
)
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
def set_seeds(seed=0):
torch.manual_seed(seed)
random.seed(seed)