github_search/pipeline.yaml

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

tasks:

  # download area grouped tasks
  - source: github_search.paperswithcode_task_areas.prepare_area_grouped_tasks
    product: "data/paperswithcode_tasks.csv"

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

  # prepare data for token2vec (modified import2vec)
  - source: github_search.data_engineering.prepare_module_corpus
    params:
      python_file_paths: ["{{python_files_csv_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_csv_path}}"]
    product: "output/papers_with_imports.csv"

  # prepare python dependency graph records
  - source: github_search.data_engineering.prepare_dependency_records
    params:
      sample_files_per_repo: 100
      add_filename_repo_label: True
      python_file_paths: ["{{python_files_csv_path}}"]
    product: "output/dependency_records.csv"

  # additional information for dependency records
  - source: github_search.data_engineering.postprocess_dependency_records
    upstream:
      - prepare_dependency_records
      - prepare_paperswithcode_with_imports_df
    params:
      use_additional_records: False
      description_mode: "{{gnn.description_mode}}" 
    product: "output/processed_dependency_records.csv"

  # make igraph from dependency records
  - source: github_search.data_engineering.make_igraph
    upstream:
      prepare_dependency_records
    product: "output/call_igraph.pkl"

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

  # get github readmes
  - source: github_search.data_engineering.make_readmes
    params:
      paperswithcode_with_tasks_path: "data/paperswithcode_with_tasks.csv"
      max_workers: 24
    upstream: prepare_paperswithcode_with_imports_df
    product: "output/papers_with_readmes.csv"

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

#  # train GraphSAGE model
#  - source: github_search.pytorch_geometric_training.run_gnn_experiment
#    upstream:
#      - postprocess_dependency_records
#      - train_python_token_fasttext
#    params:
#      fasttext_model_path: "output/python_files_fasttext_dim{{fasttext.dimension}}.bin"
#      epochs: "{{gnn.epochs}}" 
#      hidden_channels: "{{gnn.hidden_channels}}" 
#      batch_size: "{{gnn.batch_size}}" 
#      num_layers: "{{gnn.layers}}"
#      test_run: "{{test_run}}" 
#      lr: 0.001
#      model_name: "{{gnn.model_name}}"
#      use_self_connection: "{{gnn.use_self_connection}}"
#      description_mode: "{{gnn.description_mode}}" 
#    product:
#      plot_file: "output/loss_{{gnn.model_name}}_fasttext_dim{{fasttext.dimension}}_epochs{{gnn.epochs}}_dim{{gnn.hidden_channels}}_layers{{gnn.layers}}.png"
#      gnn_token_embeddings: "output/{{gnn.model_name}}_embeddings_fasttext_dim{fasttext.dimension}}_epochs{{gnn.epochs}}_dim{{gnn.hidden_channels}}_layers{{gnn.layers}}.bin"
#      model_path: "output/{{gnn.model_name}}_model_{{gnn.epochs}}_dim{{gnn.hidden_channels}}_layers{{gnn.layers}}.pth"

#  # train-test split for tasks
#  - source: github_search.matching_zsl.prepare_task_train_test_split
#    upstream:
#      - postprocess_dependency_records
#      - train_python_token_fasttext
#      - prepare_area_grouped_tasks
#    product:
#      train: "output/tasks_train.csv"
#      test: "output/tasks_test.csv"

  # task embeddings
  # prepares data for task_embeddings_app streamlit app
  # depends on sentence_embeddings_main
  - source: github_search.embedding_utils.prepare_reduced_embeddings
    upstream:
      - train_abstract_readme_w2v
      - train_python_token_fasttext
    params:
      rnn_model_path: "output/sbert/sru2x256_epoch350"
      codebert_model_path: "output/sbert/codebert15"
    product: "output/reduced_features.pkl"

  # for each repo select files that are most similar
  # similarity is cosine similarity in provided model
  - source: github_search.python_code_analysis.select_repo_files
    upstream:
      - make_readmes 
    params:
      similar_col: "readme"
      bow_vectorizer_class: "CountVectorizer"
      files_per_repo: 10
      python_files_path: "{{python_files_csv_path}}"
    product: "output/selected_python_files.feather"


  - source: github_search.python_code_analysis.extract_python_comments
    upstream:
      - select_repo_files
    params:
      line_neighbor_threshold: 2
    product:
      "output/selected_python_files_comments.feather"

  - source: github_search.python_code_analysis.extract_python_imports
    upstream:
      - select_repo_files
    product:
      "output/selected_python_files_imports.feather"

  - source: github_search.seq2seq_utils.prepare_path_docid_seq2seq_df
    upstream:
      - select_repo_files
      - make_readmes
    product:
      "output/path_docid_seq2seq.csv"

  - source: github_search.seq2seq_utils.prepare_seq2seq_dataset
    name: prepare_path_docid_seq2seq_dataset
    upstream:
      - prepare_path_docid_seq2seq_df
    params:
      base_model: "{{transformer.T5_model}}" 
      max_source_length: "{{transformer.max_source_length}}" 
      max_target_length: "{{transformer.max_target_length}}" 
    product:
      "output/seq2seq_hf_dataset"