Update app.py

app.py

@@ -1,15 +1,27 @@
+import os
 import torch
+import logging
 from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer
 from datasets import load_dataset
 
+# Set verbose logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# Set a writable cache directory
+os.environ["HF_HOME"] = "/app/hf_cache"
+os.environ["TRANSFORMERS_CACHE"] = "/app/hf_cache"
+
 # Load dataset
-print("Loading dataset...")
+logger.info("Loading dataset...")
 ds = load_dataset("facebook/natural_reasoning")
+logger.info(f"Dataset loaded successfully! Dataset info:\n{ds}")
 
 # Load tokenizer
-print("Loading 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):
@@ -17,14 +29,17 @@ def preprocess_function(examples):
     return tokenizer(input_texts, truncation=True, padding="max_length", max_length=512)
 
 # Tokenize dataset
-print("Tokenizing dataset...")
+logger.info("Tokenizing dataset...")
 tokenized_datasets = ds.map(preprocess_function, batched=True)
+logger.info("Dataset tokenized successfully!")
 
 # Load model
-print("Loading model...")
+logger.info("Loading model...")
 model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True)
+logger.info("Model loaded successfully!")
 
 # Training arguments
+logger.info("Setting up training arguments...")
 training_args = TrainingArguments(
     output_dir="./results",
     evaluation_strategy="epoch",
@@ -35,10 +50,14 @@ training_args = TrainingArguments(
     weight_decay=0.01,
     logging_dir="./logs",
     logging_steps=10,
-    push_to_hub=True  # Upload trained model to Hugging Face Hub
+    push_to_hub=True,  # Upload trained model to Hugging Face Hub
+    report_to="none",  # Prevents sending logs to external services
+    logging_first_step=True
 )
+logger.info("Training arguments set!")
 
 # Trainer
+logger.info("Initializing Trainer...")
 trainer = Trainer(
     model=model,
     args=training_args,
@@ -46,12 +65,14 @@ trainer = Trainer(
     eval_dataset=tokenized_datasets["test"],
     tokenizer=tokenizer
 )
+logger.info("Trainer initialized!")
 
 # Start training
-print("Starting training...")
+logger.info("Starting training...")
 trainer.train()
+logger.info("Training completed!")
 
 # Push trained model to Hugging Face Hub
-print("Pushing model to Hub...")
+logger.info("Pushing trained model to Hugging Face Hub...")
 trainer.push_to_hub()
-print("Training complete!")
+logger.info("Model push completed! Training process finished successfully.")