Enable more config inputs to gqe. (#11)

Signed-off-by: Alex McCaskey <[email protected]>

Author: amccaskey

cudaqlib/operators/chemistry/drivers/pyscf/ExternalPySCFDriver.cpp

@@ -88,7 +88,7 @@ class external_pyscf
     // Import all the data we need from the execution.
     std::ifstream f(metadataFile);
     auto metadata = nlohmann::json::parse(f);
-    printf("TEST\n%s\n", metadata.dump(4).c_str());
+    
     // Get the energy, num orbitals, and num qubits
     std::unordered_map<std::string, double> energies;
     for (auto &[energyName, E] : metadata["energies"].items())

python/cudaqlib/algorithms/gqe.py

@@ -223,11 +223,16 @@ def __internal_run_gqe(temperature_scheduler: TemperatureScheduler,
     fabric.log('circuit', json.dumps(min_indices))
     return min_energy, min_indices
 
-
 def gqe(cost, pool, config=None, **kwargs):
-    cfg = get_default_config() if config == None else config
+    cfg = get_default_config() 
+    if config == None: 
+        [setattr(cfg, a, kwargs[a]) for a in dir(cfg) if not a.startswith('_') and a in kwargs]
+    else:
+        cfg = config
+
+    # Don't let someone override the vocab_size
     cfg.vocab_size = len(pool)
-    if 'max_iters' in kwargs: cfg.max_iters = kwargs['max_iters']
+
     model = Transformer(
         cfg, cost, loss='exp') if 'model' not in kwargs else kwargs['model']
     optimizer = torch.optim.AdamW(

python/cudaqlib/algorithms/transformer.py

@@ -4,6 +4,7 @@
 from lightning import LightningModule
 from .loss import ExpLogitMatching, GFlowLogitMatching
 
+
 def get_device():
     if torch.cuda.is_available():
         return 'cuda'
@@ -20,16 +21,18 @@ def __init__(self, **kwargs):
 
 class Transformer(LightningModule):
 
-    def __init__(self, cfg, cost,  loss="exp"):
+    def __init__(self, cfg, cost, loss="exp"):
         super().__init__()
         self._label = 'label_stand_in'
         self.cfg = cfg
         gpt2cfg = GPT2Config(
-            **{k: cfg[k] for k in GPT2Config().to_dict().keys() & cfg.keys()})
+            **{k: cfg[k]
+               for k in GPT2Config().to_dict().keys() & cfg.keys()})
         if cfg.small:
-            gpt2cfg = SmallConfig(
-                **
-                {k: cfg[k] for k in GPT2Config().to_dict().keys() & cfg.keys()})
+            gpt2cfg = SmallConfig(**{
+                k: cfg[k]
+                for k in GPT2Config().to_dict().keys() & cfg.keys()
+            })
         self.transformer = GPT2LMHeadModel(gpt2cfg).to(get_device())
         self.ngates = cfg.ngates
         self.num_samples = cfg.num_samples
@@ -57,22 +60,31 @@ def set_cost(self, cost):
 
     def gather(self, idx, logits_base):
         b_size = idx.shape[0]
-        return torch.gather(logits_base, 2, idx.reshape(b_size, -1,
-                                                        1)).reshape(b_size, -1)
+        return torch.gather(logits_base, 2,
+                            idx.reshape(b_size, -1, 1)).reshape(b_size, -1)
 
     def computeCost(self, idx_output, pool, **kwargs):
-        return torch.tensor([self._cost([pool[i] for i in row]) for row in idx_output],
-                            dtype=torch.float)
+        return torch.tensor(
+            [self._cost([pool[i] for i in row]) for row in idx_output],
+            dtype=torch.float)
 
-    def train_step(self, pool, indices=None, energies=None, numQPUs=None, comm=None):
+    def train_step(self,
+                   pool,
+                   indices=None,
+                   energies=None,
+                   numQPUs=None,
+                   comm=None):
         log_values = {}
         if energies is not None:
             assert indices is not None
             idx_output = indices[:, 1:]
             logits_base = self.generate_logits(idx_output)
         else:
             idx_output, logits_base = self.generate()
-            energies = self.computeCost(idx_output, pool, numQPUs=numQPUs, comm=comm)
+            energies = self.computeCost(idx_output,
+                                        pool,
+                                        numQPUs=numQPUs,
+                                        comm=comm)
         logits_tensor = self.gather(idx_output, logits_base)
         allLogits = logits_tensor
 

tools/chemistry/cudaq-pyscf.py

@@ -94,7 +94,6 @@ def iter_namespace(ns_pkg):
         for (k, v) in vars(args).items() if k not in filterArgs
     }
     res = hamiltonianGenerator.generate(args.xyz, args.basis, **filteredArgs)
-    print(res)
 
     exit(0)