Update app.py

app.py

@@ -1,44 +1,57 @@
-from datasets import load_dataset
-from transformers import AutoModelForSequenceClassification, TrainingArguments, Trainer, AutoTokenizer
 import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer
+from datasets import load_dataset
 
-# Load Dataset
-dataset = load_dataset("yelp_review_full")  # Example dataset
+# Load dataset
+print("Loading dataset...")
+ds = load_dataset("facebook/natural_reasoning")
 
-# Load Pretrained Model & Tokenizer
-model_name = "bert-base-uncased"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=5)
+# Load tokenizer
+print("Loading tokenizer...")
+model_name = "deepseek-ai/DeepSeek-R1"
+tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
 
-# Tokenize Dataset
+# Tokenization function
 def preprocess_function(examples):
-    return tokenizer(examples["text"], truncation=True, padding="max_length", max_length=512)
+    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
+print("Tokenizing dataset...")
+tokenized_datasets = ds.map(preprocess_function, batched=True)
 
-encoded_dataset = dataset.map(preprocess_function, batched=True)
+# Load model
+print("Loading model...")
+model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True)
 
-# Training Arguments
+# Training arguments
 training_args = TrainingArguments(
     output_dir="./results",
     evaluation_strategy="epoch",
     save_strategy="epoch",
-    per_device_train_batch_size=8,
-    per_device_eval_batch_size=8,
+    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,
-    push_to_hub=True  # Push trained model back to Hugging Face
+    logging_dir="./logs",
+    logging_steps=10,
+    push_to_hub=True  # Upload trained model to Hugging Face Hub
 )
 
-# Define Trainer
+# Trainer
 trainer = Trainer(
     model=model,
     args=training_args,
-    train_dataset=encoded_dataset["train"],
-    eval_dataset=encoded_dataset["test"],
+    train_dataset=tokenized_datasets["train"],
+    eval_dataset=tokenized_datasets["test"],
     tokenizer=tokenizer
 )
 
-# Train the Model
+# Start training
+print("Starting training...")
 trainer.train()
 
-# Save & Push to Hub
+# Push trained model to Hugging Face Hub
+print("Pushing model to Hub...")
 trainer.push_to_hub()
+print("Training complete!")