small fixes

Author: durant42040

lean_dojo_v2/__init__.py

@@ -1,17 +1,3 @@
 
 __version__ = "1.0.0"
 __author__ = "LeanDojo-v2 Contributors"
-
-# Import main components for easy access
-from .agent import BaseAgent, HFAgent, LeanAgent
-from .prover import BaseProver, ExternalProver, HFProver, RetrievalProver
-
-__all__ = [
-    "BaseAgent",
-    "HFAgent",
-    "LeanAgent",
-    "BaseProver",
-    "HFProver",
-    "RetrievalProver",
-    "ExternalProver",
-]

lean_dojo_v2/lean_progress/train_steps_model.py

@@ -73,7 +73,8 @@ def _tokenize(batch):
             truncation=True,
             max_length=max_length,
         )
-        encoded["labels"] = batch["steps_remaining"]
+        # encoded["labels"] = batch["steps_remaining"]
+        encoded["labels"] = [float(x) for x in batch["steps_remaining"]]
         return encoded
 
     return _tokenize
@@ -122,7 +123,7 @@ def main() -> None:
         per_device_eval_batch_size=args.batch_size,
         num_train_epochs=args.epochs,
         weight_decay=0.01,
-        evaluation_strategy="epoch",
+        eval_strategy="epoch",
         save_strategy="epoch",
         load_best_model_at_end=True,
         metric_for_best_model="mse",