BriLLM

/

BriLLM0.5

like 34

It just takes a long time, you can add a progress bar to this part to wait for it to complete.

        from tqdm import tqdm
        total = len(vocab.edge_dict)
        for s_idx, s in tqdm(enumerate(vocab.edge_dict), total=total, desc="edge_dict", unit="item", ncols=100, mininterval=0.1, bar_format='{l_bar}{bar}| {n_fmt}/{total_fmt} [{percentage:3.1f}%]'):
            for t_idx, t in enumerate(vocab.edge_dict[s]):
                if self.zero_freq_edges is not None and t in self.zero_freq_edges[s]:
                    # Use shared parameters
                    self.weight_indices[(s_idx, t_idx)] = self.shared_param_idx
                else:
                    self.weight_indices[(s_idx, t_idx)] = current_idx
                    current_idx += 1

replace the Nested loops in model.prepare_network to this.