Update app.py

app.py

@@ -4,17 +4,20 @@ 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__)
 
-# Set a writable cache directory
-os.environ["HF_HOME"] = "/app/hf_cache"
-os.environ["TRANSFORMERS_CACHE"] = "/app/hf_cache"
-
-# Load dataset
 logger.info("Loading dataset...")
-ds = load_dataset("facebook/natural_reasoning")
+ds = load_dataset("facebook/natural_reasoning")  # Replace with your dataset
 logger.info(f"Dataset loaded successfully! Dataset info:\n{ds}")
 
 # Load tokenizer
@@ -39,7 +42,6 @@ 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",
@@ -50,14 +52,12 @@ training_args = TrainingArguments(
     weight_decay=0.01,
     logging_dir="./logs",
     logging_steps=10,
-    push_to_hub=True,  # Upload trained model to Hugging Face Hub
-    report_to="none",  # Prevents sending logs to external services
+    push_to_hub=True,
+    report_to="none",
     logging_first_step=True
 )
-logger.info("Training arguments set!")
 
 # Trainer
-logger.info("Initializing Trainer...")
 trainer = Trainer(
     model=model,
     args=training_args,
@@ -65,7 +65,6 @@ trainer = Trainer(
     eval_dataset=tokenized_datasets["test"],
     tokenizer=tokenizer
 )
-logger.info("Trainer initialized!")
 
 # Start training
 logger.info("Starting training...")