Spaces:
Runtime error
Runtime error
import os | |
import torch | |
import logging | |
from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer | |
from datasets import load_dataset | |
# ✅ Set a writable cache directory inside the container | |
os.environ["HF_HOME"] = "/app/hf_cache" | |
os.environ["TRANSFORMERS_CACHE"] = "/app/hf_cache" | |
os.environ["HF_DATASETS_CACHE"] = "/app/hf_cache" | |
# Ensure cache directory exists | |
os.makedirs("/app/hf_cache", exist_ok=True) | |
# Set verbose logging | |
logging.basicConfig(level=logging.INFO) | |
logger = logging.getLogger(__name__) | |
logger.info("Loading dataset...") | |
ds = load_dataset("facebook/natural_reasoning") # Replace with your dataset | |
logger.info(f"Dataset loaded successfully! Dataset info:\n{ds}") | |
# Load tokenizer | |
logger.info("Loading tokenizer...") | |
model_name = "deepseek-ai/DeepSeek-R1" | |
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True) | |
logger.info("Tokenizer loaded successfully!") | |
# Tokenization function | |
def preprocess_function(examples): | |
input_texts = [f"Q: {q} A: {a}" for q, a in zip(examples["question"], examples["reference_answer"])] | |
return tokenizer(input_texts, truncation=True, padding="max_length", max_length=512) | |
# Tokenize dataset | |
logger.info("Tokenizing dataset...") | |
tokenized_datasets = ds.map(preprocess_function, batched=True) | |
logger.info("Dataset tokenized successfully!") | |
# Load model | |
logger.info("Loading model...") | |
model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True) | |
logger.info("Model loaded successfully!") | |
# Training arguments | |
training_args = TrainingArguments( | |
output_dir="./results", | |
evaluation_strategy="epoch", | |
save_strategy="epoch", | |
per_device_train_batch_size=4, # Adjust based on available RAM | |
per_device_eval_batch_size=4, | |
num_train_epochs=3, | |
weight_decay=0.01, | |
logging_dir="./logs", | |
logging_steps=10, | |
push_to_hub=True, | |
report_to="none", | |
logging_first_step=True | |
) | |
# Trainer | |
trainer = Trainer( | |
model=model, | |
args=training_args, | |
train_dataset=tokenized_datasets["train"], | |
eval_dataset=tokenized_datasets["test"], | |
tokenizer=tokenizer | |
) | |
# Start training | |
logger.info("Starting training...") | |
trainer.train() | |
logger.info("Training completed!") | |
# Push trained model to Hugging Face Hub | |
logger.info("Pushing trained model to Hugging Face Hub...") | |
trainer.push_to_hub() | |
logger.info("Model push completed! Training process finished successfully.") | |