github_search/pipeline.yaml

executor:
  dotted_path: ploomber.executors.Serial
  build_in_subprocess: False

meta:
  extract_upstream: False 
  extract_product: False
  jupyter_hot_reload: True
  jupyter_functions_as_notebooks: True

tasks:

  # download area grouped tasks
  - source: github_search.papers_with_code.paperswithcode_task_areas.prepare_area_grouped_tasks
    product: "data/paperswithcode_tasks.csv"
  
  - source: github_search.papers_with_code.ploomber.prepare_paperswithcode_df
    name: pwc_data.prepare_raw_paperswithcode_df
    params:
      paperswithcode_filename: "data/links-between-papers-and-code.json.gz"
      papers_filename: "data/papers-with-abstracts.json.gz"
    product:
      paperswithcode_path: "output/raw_paperswithcode_df.csv"

  - source: github_search.papers_with_code.ploomber.prepare_filtered_paperswithcode_df
    name: pwc_data.prepare_final_paperswithcode_df
    upstream:
      - pwc_data.prepare_raw_paperswithcode_df
    params:
      min_task_count: 10
    product:
      paperswithcode_path: "{{paperswithcode_path}}"
      task_counts_path: "output/task_counts.csv"

  # READMEs
  # get github readmes
  - source: github_search.elixir_runner.download_readmes_pb
    name: pwc_data.download_readmes
    upstream:
      - pwc_data.prepare_final_paperswithcode_df
    product: "output/paperswithcode_readmes.json"

  - source: github_search.papers_with_code.ploomber.prepare_paperswithcode_with_readmes_pb
    name: pwc_data.prepare_paperswithcode_with_readmes
    upstream:
      - pwc_data.prepare_final_paperswithcode_df
      - pwc_data.download_readmes
    product: "output/paperswithcode_with_readmes.json.gz"

  # train-test split for tasks
  # tasks are stratified by paperswithcode area
  - source: github_search.train_test_split.prepare_task_train_test_split
    upstream:
      - prepare_area_grouped_tasks
      - pwc_data.prepare_final_paperswithcode_df
    params:
     test_size: 1
    product:
      train: "output/tasks_train.csv"
      test: "output/tasks_test.csv"


  #######################
  # code2doc
  #######################
  - source: github_search.pipelines.steps.code2doc_prepare_data_pb
    name: code2doc.prepare_data
    params:
      repos_df_path: "output/paperswithcode_with_readmes.json.gz"
      python_code_path: "output/repo_selected_files.parquet"
    product:
      repos_df_path: "output/repos_with_all_data.jsonl"
      # for some reason there are errors in parquet so we'll save it to feather
      selected_python_code_path: "output/selected_python_code.feather"

  - source: github_search.pipelines.steps.create_repos_sample_pb
    name: code2doc.create_repo_sample
    upstream:
      - code2doc.prepare_data
    params:
      min_task_size: 5
      n_repos_per_task: 10
      max_task_size: 500
      max_random_baseline_score: 0.5
    product:
      sampled_repos: "output/code2doc/sample_per_task_5_repos/sampled_repos5.jsonl"

  # to run this step serve ollama on the appropriate port
  - source: github_search.pipelines.steps.generate_code2doc_readmes_pb
    name: code2doc.generate_readmes
    upstream:
      - code2doc.create_repo_sample
      - code2doc.prepare_data
    params:
      lm_model_name: "codellama"
      lm_base_url: "http://localhost:11430"
      small_lm_base_url: "http://localhost:11431"
      files_per_repo: 10
    product:
      "output/code2doc/sample_per_task_5_repos/codellama_generated_readmes5.jsonl"
  # the stack
  - source: github_search.the_stack.prepare_the_stack_files
    name:
      the_stack.prepare_files
    params:
      paperswithcode_path: "data/paperswithcode_with_tasks.csv"
      delete_temporary_files: True
    product: "data/the_stack_paperswithcode_repos"

  - source: github_search.the_stack.prepare_the_stack_df
    name:
      the_stack.prepare_df
    upstream:
      - the_stack.prepare_files
    product: "output/the_stack_paperswithcode_files.parquet"

  # prepare data for similarity learning from paperswithcode
  - source: github_search.sentence_embeddings.datasets.prepare_paperswithcode_data
    name: sentence_embeddings.prepare_paperswithcode_data
    product:
      datasets: "data/datasets.json.gz"
      methods: "data/methods.json.gz"

  # prepare data for similarity learning from dbpedia 
  - source: github_search.sentence_embeddings.datasets.prepare_dbpedia_machine_learning_data
    name: sentence_embeddings.prepare_dbpedia_data
    product: "data/dbpedia_ml_records.csv" 

  - source: github_search.sentence_embeddings.datasets.prepare_data
    name: sentence_embeddings.prepare_data
    upstream:
      - sentence_embeddings.prepare_dbpedia_data
      - sentence_embeddings.prepare_paperswithcode_data
    product: "data/sentence_similarity_data.csv"

  #
  - source: github_search.train_test_split.prepare_repo_train_test_split
    upstream:
      - pwc_data.prepare_paperswithcode_with_readmes
      - prepare_task_train_test_split
    product:
      train: "output/repos_train.json"
      test: "output/repos_test.json"

  # extract python tokens for BoW baseline
  - source: github_search.bow_baseline.extract_python_tokens
    product: "output/python_files_with_tokens_df.feather"

  # 
  - source: github_search.bow_baseline.prepare_bow_retrieval_evaluation_results
    name: prepare_bow_retrieval_evaluation_results_readme
    upstream:
      - prepare_task_train_test_split
    params:
      index: python_tokenized_files 
    product: "output/python_files_retrieval_results.csv"
  # 
  - source: github_search.bow_baseline.prepare_bow_retrieval_evaluation_results
    name: prepare_bow_retrieval_evaluation_results_python_files
    upstream:
      - prepare_task_train_test_split
    params:
      index: project_readmes 
    product: "output/readme_retrieval_results.csv"

  # run word2vec on natural language data
  - source: github_search.word2vec.train_abstract_readme_w2v
    upstream:
      - pwc_data.prepare_paperswithcode_with_readmes
    params:
      embedding_dim: "{{word2vec.dimension}}"
      epochs: "{{word2vec.epochs}}"
    product:
      binary: "output/abstract_readme_w2v{{word2vec.dimension}}.bin"
      txt: "output/abstract_readme_w2v{{word2vec.dimension}}.txt"

  # run word2vec on code 
  - source: github_search.word2vec.train_python_code_w2v
    params:
      python_file_path: "{{python_files_path}}"
      embedding_dim: "{{word2vec.dimension}}"
    product:
      binary: "output/python_code_w2v{{word2vec.dimension}}.bin"
      txt: "output/python_code_w2v{{word2vec.dimension}}.txt"

  #
  # imports
  #
  - source: github_search.imports.prepare_data.prepare_file_imports
    name: imports.prepare_file_imports
    params:
      python_files_path: "{{python_files_path}}"
    product: "output/python_file_imports.feather"

  - source: github_search.imports.training.train_import_word2vec
    name: imports.train_w2v
    upstream:
      - imports.prepare_file_imports
    params:
      embedding_dim: "{{word2vec.dimension}}"
      epochs: "{{word2vec.epochs}}"
    product:
      binary: "output/imports_w2v{{word2vec.dimension}}.bin"
      txt: "output/imports_w2v{{word2vec.dimension}}.txt"

  - source: github_search.imports.training.train_import_rnn_file_similarity_model
    name: imports.train_rnn
    upstream:
      - imports.prepare_file_imports
      - imports.train_w2v
    params:
      epochs: 2
      batch_size: 256 
      rnn_config: "{{rnn_config}}"
    product: "output/models/import_lstm"

  #
  # sentence embeddings
  #

  # make word2vec aggregator model
  - source: github_search.sentence_embeddings.models.prepare_word2vec_sentence_embedding_model
    name: sentence_embeddings.prepare_w2v_model
    upstream:
      - train_abstract_readme_w2v
    product:
      "output/abstract_readme_embedder"

  # prepare data for token2vec (modified import2vec)
  - source: github_search.data_engineering.prepare_module_corpus
    params:
      python_file_paths: ["{{python_files_path}}"]
    product: "output/module_corpus.csv"

  # train token2vec model
  - source: github_search.token2vec.train_token2vec
    upstream:
      - prepare_module_corpus
    params:
      n_iterations: 10000
      n_positive_imports: 32
      embedding_dim: "{{word2vec.dimension}}" 
    product:
      model_path: "output/import2vec_module_vectors{{word2vec.dimension}}.bin"

  # prepare paper dataset with imports extracted per-project
  - source: github_search.data_engineering.prepare_paperswithcode_with_imports_df
    upstream:
      - prepare_module_corpus
    params:
      python_file_paths: ["{{python_files_path}}"]
    product: "output/papers_with_imports.csv"

  # prepare python dependency graph records
  - source: github_search.data_engineering.prepare_dependency_records
    name:
      dependency_graph.prepare_records
    params:
      sample_files_per_repo: 1000
      add_repo_col: True
      use_basename: False
      python_file_path: "{{python_files_path}}"
      excluded_prefix: "venv"
    product: "output/dependency_records.feather"

  # additional information for dependency records
  - source: github_search.data_engineering.postprocess_dependency_records
    name:
      dependency_graph.postprocess_records
    upstream:
      - dependency_graph.prepare_records
      - prepare_paperswithcode_with_imports_df
    params:
      use_additional_records: False
      description_mode: False
    product: "output/processed_dependency_records.feather"

  #
  #
  # GRAPHS
  #
  #

  # Records from Neo4J
  - source: github_search.neo4j_graph.prepare_neo4j_dependency_records
    upstream:
      - prepare_repo_train_test_split
    params:
      graph_dependencies_path: "output/dependency_records/repo_dependencies_articlerank.json"
      id_col: repo
      rel_col: edge_type
    product:
      train: "output/dependency_records/graph_dependencies_train.json"
      test: "output/dependency_records/graph_dependencies_test.json"

  # extract python function df
  # f has columns
  # ['repo_name', 'path', 'function_name', 'function_code']
  - source: github_search.data_engineering.prepare_function_code_df
    params:
      python_file_path: "{{python_files_path}}"
      max_depth: 10
      n_cores: 4
    product:
      "output/python_functions.feather"

  - source: github_search.data_engineering.prepare_function_signatures_df
    params:
      python_file_path: "{{python_files_path}}"
      n_cores: 4
    product:
      "output/python_signatures.parquet"


  # train FastText model on Python files
  - source: github_search.data_engineering.train_python_token_fasttext
    params:
      python_file_path: "{{python_files_path}}"
      dim: "{{fasttext.dimension}}" 
      epoch: "{{fasttext.epochs}}"
      n_cores: 16 
    product:
      "output/python_files_fasttext_dim{{fasttext.dimension}}.bin"

  #
  - source: github_search.summarization.prepare_function_df_with_summarized_code
    params:
      transformer_model_name: "{{summarization.transformer_model_name}}"
    upstream: prepare_function_code_df
    product:
      "output/python_files_descriptions_{{summarization.transformer_model_name}}.feather"


  - source: github_search.graphs.prepare_graph.prepare_from_dependency_records
    name: graph.prepare_from_dependency_records
    upstream:
      - dependency_graph.prepare_records
    params:
      used_edges:
        - "repo-file"
    product: "output/dependency_records_igraph.pkl"

  - source: github_search.graphs.prepare_graph.prepare_from_function_code
    name: graph.prepare_from_function_code
    upstream:
      - prepare_function_code_df
    product: "output/function_code_igraph.pkl"

  # prepare graph list 
  - source: github_search.graphs.data_preparation.prepare_dataset_with_transformer
    name: gnn.prepare_dataset_with_transformer
    params:
      sentence_transformer_model_name: "{{gnn.sentence_transformer_model_name}}"
      batch_size: 128
      paperswithcode_path: "{{paperswithcode_path}}"
    upstream:
      - prepare_area_grouped_tasks
      - graph.prepare_from_function_code
    product: "output/graph_list.pkl"