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)