Update app.py

app.py

@@ -16,9 +16,15 @@ import spaces
 # Use dotenv to load the environment variables
 load_dotenv()
 
-# Get HF token from environment variable
+# Get environment variables
 HF_TOKEN = os.getenv("HF_TOKEN_TEXT")
-print(HF_TOKEN)
+DATASET_NAME = os.getenv("DATASET_NAME", "QuotaClimat/frugalaichallenge-text-train")  # Default to public dataset
+TEST_SIZE = float(os.getenv("TEST_SIZE", "0.2"))  # Default to 20% test size
+TEST_SEED = int(os.getenv("TEST_SEED", "42"))  # Default seed for reproducibility
+
+print(f"Using dataset: {DATASET_NAME}")
+print(f"Test split size: {TEST_SIZE}")
+
 if not HF_TOKEN:
     print("Warning: HF_TOKEN not found in environment variables. Submissions will not work.")
 
@@ -162,13 +168,13 @@ LABEL_MAPPING = {
 
 # Load and prepare the dataset
 print("Loading dataset...")
-dataset = load_dataset("QuotaClimat/frugalaichallenge-text-train")
+dataset = load_dataset(DATASET_NAME)
 
 # Convert string labels to integers
 dataset = dataset.map(lambda x: {"label": LABEL_MAPPING[x["label"]]})
 
 # Split dataset
-train_test = dataset["train"].train_test_split(test_size=0.2, seed=42)
+train_test = dataset["train"].train_test_split(test_size=TEST_SIZE, seed=TEST_SEED)
 train_dataset = train_test["train"]
 test_dataset = train_test["test"]