Update PolyAgent/orchestrator.py

PolyAgent/orchestrator.py

@@ -1,5 +1,5 @@
 """
-PolyAgent Orchestrator (5M)
+PolyAgent Orchestrator
 ===========================
 
 This file provides a modular orchestrator that:
@@ -26,11 +26,11 @@ import numpy as np
 import torch
 import torch.nn as nn
 
-# HF Transformers (for SELFIES-TED decoder as in G2.py)
+# HF Transformers (for SELFIES-TED decoder)
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 from transformers.modeling_outputs import BaseModelOutput
 
-# Optional imports for web fetching
+# Imports for web fetching
 try:
     import requests
     from bs4 import BeautifulSoup
@@ -38,7 +38,7 @@ except Exception:
     requests = None
     BeautifulSoup = None
 
-# Optional imports for visuals
+# Imports for visuals
 try:
     from rdkit import Chem
     from rdkit.Chem import Draw
@@ -51,7 +51,7 @@ try:
 except Exception:
     cm = None
 
-# Optional: joblib + sentencepiece for 5M generator artifacts
+# joblib + sentencepiece for 5M generator artifacts
 try:
     import joblib
 except Exception:
@@ -62,7 +62,7 @@ try:
 except Exception:
     spm = None
 
-# Optional: selfies (for SELFIES→SMILES/PSMILES conversion)
+# selfies (for SELFIES→SMILES/PSMILES conversion)
 try:
     import selfies as sf
 except Exception:
@@ -85,13 +85,11 @@ class PathsConfig:
     """
 
     def __init__(self):
-        # 1) HF model repo where you uploaded the staged bundle
-        #    Example: "kaurm43/PolyFusionAgent-weights-5m" (change to your real repo_id)
+        # 1) HF model repo 
         self.hf_repo_id = os.getenv("POLYFUSION_WEIGHTS_REPO", "kaurm43/polyfusionagent-weights")
         self.hf_repo_type = os.getenv("POLYFUSION_WEIGHTS_REPO_TYPE", "model")  # usually "model"
 
         # 2) Where to store downloaded files
-        #    Prefer /data on Spaces with persistent storage; else use a cache folder.
         default_root = "/data/polyfusion_cache" if os.path.isdir("/data") else os.path.expanduser("~/.cache/polyfusion_cache")
         self.local_weights_root = os.getenv("POLYFUSION_WEIGHTS_DIR", default_root)
 
@@ -99,7 +97,6 @@ class PathsConfig:
         self.hf_token = os.getenv("HF_TOKEN", None)
 
         # 4) Download (cached) + get local folder path.
-        #    allow_patterns keeps download smaller/faster (only pull what orchestrator needs).
         allow = [
             "tokenizer_spm_5m/**",
             "polyfusion_cl_5m/**",
@@ -117,8 +114,7 @@ class PathsConfig:
             allow_patterns=allow,
         )
 
-        # 5) Map to the exact files your existing code expects
-        #    (Only path wiring changes; no behavior changes elsewhere.)
+        # 5) Map to the necessary files 
         self.cl_weights_path = os.path.join(self._weights_dir, "polyfusion_cl_5m", "pytorch_model.bin")
 
         # If your Space also includes a local Chroma DB folder in the Space repo,
@@ -131,7 +127,7 @@ class PathsConfig:
         self.downstream_bestweights_5m_dir = os.path.join(self._weights_dir, "downstream_heads_5m")
         self.inverse_design_5m_dir = os.path.join(self._weights_dir, "inverse_design_5m")
 
-        # 6) Optional: sanity-check required files (fail early with a clear message)
+        # 6) Optional: sanity-check required files 
         self._assert_exists(self.cl_weights_path, "CL weights")
         self._assert_exists(self.spm_model_path, "SentencePiece model")
         self._assert_exists(self.spm_vocab_path, "SentencePiece vocab")
@@ -601,7 +597,7 @@ def _assign_tool_tags_to_report(report: Dict[str, Any]) -> Dict[str, Any]:
 
 def _render_tool_outputs_verbatim_md(report: Dict[str, Any]) -> str:
     """
-    Render tool outputs as verbatim JSON blocks (no content rewriting).
+    Render tool outputs as verbatim JSON blocks.
     """
     if not isinstance(report, dict):
         return ""