Class Pipeline

src/pipeline.py:23–119

class Pipeline:
    def __init__(self, config: Config):
        self.config = config
        self._obj_map = instantiate_class_from_config(config)
        self.pipeline_settings = self._obj_map["pipeline_settings"][0]
        self.semantic_graph_context_generator = [
            self._obj_map["semantic_graph_context_generator"][i]
            for i in range(len(self._obj_map["semantic_graph_context_generator"]))
        ]
        self.gnn_heads = [self._obj_map["gnn_heads"][i] for i in range(len(self._obj_map["gnn_heads"]))]
        self.topic_model = self._obj_map["topic_model"][0]

    @classmethod
    def from_config(cls, config: Config):
        return cls(config)

    @classmethod
    def from_yaml(cls, yaml_path):
        config = Config.from_yaml(yaml_path)
        return cls(config)

    def find_files_with_substring(self, root_dir, substring):
        for dirpath, dirnames, filenames in os.walk(root_dir):
            for filename in filenames:
                if substring in filename:
                    yield os.path.join(dirpath, filename)

    def load_jsonl(self, filepaths):
        res = []
        for filepath in filepaths:
            with open(filepath, "r") as f:
                for line in f:
                    res.append(json.loads(line))
        sents = []
        for r in res:
            messages = r["conversation_history"]
            reply = r["assistant_reply"]
            sents.append(reply)
            sents.append(messages[-2]["content"])
        data = pd.DataFrame(sents, columns=["query"])
        data["_id"] = data.index
        return data

    def run(self, git_repo: str, repo_name: str) -> Tuple[Tuple[np.ndarray, np.ndarray], dict]:
        """
        Run the pipeline."""
        # replace with lg.info
        lg.info("Running pipeline...")
        lg.info("Fetching repo...")
        repo_folder, context_folder = clone_and_create_context_folder(git_repo, repo_name)
        lg.info("Generating semantic graph context...")

        semantic_graph_context = [
            context_generator.decompose_repo(
                repo_folder,
                repo_name,
                context_folder,
                skip_graph_generation=self.pipeline_settings.config["skip_graph_creation"],