remove non pyfiles

Author: vpj

python_autocomplete/remove_non_source_files.py

@@ -0,0 +1,28 @@
+from pathlib import Path
+from typing import Set
+
+from labml import lab
+
+
+def remove_files(path: Path, keep: Set[str]):
+    """
+    Remove files
+    """
+
+    for p in path.iterdir():
+        if p.is_symlink():
+            p.unlink()
+            continue
+        if p.is_dir():
+            remove_files(p, keep)
+        else:
+            if p.suffix not in keep:
+                p.unlink()
+
+
+def main():
+    remove_files(lab.get_data_path() / 'source', {'.py'})
+
+
+if __name__ == '__main__':
+    main()

readme.md

@@ -11,12 +11,13 @@ This repo trains deep learning models on source code.
  [PyTorch awesome list](https://github.com/bharathgs/Awesome-pytorch-list).
 4. Run `python_autocomplete/extract_downloads.py` to extract the downloaded zip files to `data/source`.
  You can directly copy any python code to `data/source` to train on them.
-5. Run `create_dataset.py` to collect all python files.
+5. Run `python_autocomplete/remove_non_source_files.py` to all files except `.py` files.
+6. Run `create_dataset.py` to collect all python files.
  The collected code will be written to `data/train.py` and, `data/eval.py`.
-6. Run `train.py` to train the model.
+7. Run `train.py` to train the model.
  *Try changing hyper-parameters like model dimensions and number of layers*.
-7. Run `evaluate.py` to evaluate the model.
-8. Enjoy!
+8. Run `evaluate.py` to evaluate the model.
+9. Enjoy!
 
 If you have any questions please open an issue on Github.