new: Added support of configuration file.

Author: cedricbonhomme

.gitignore

@@ -34,5 +34,6 @@ eproject.cfg
 *.log
 
 emissions.csv
+vulntrain/conf.py
 vulnerability/
 

CHANGELOG.md

@@ -1,5 +1,10 @@
 # Changelog
 
+## Release 0.5.0 (2025-02-21)
+
+Added support of configuration file.
+
+
 ## Release 0.4.0 (2025-02-21)
 
 The dataset generation step now uses data from GitHub Advisories,

pyproject.toml

@@ -5,7 +5,7 @@ build-backend = "poetry.core.masonry.api"
 
 [project]
 name = "VulnTrain"
-version = "0.4.0"
+version = "0.5.0"
 description = "Generate datasets amd models based on vulnerabilities descriptions from Vulnerability-Lookup."
 authors = [
     {name = "Cédric Bonhomme",email = "[email protected]"}

vulntrain/conf.py.sample

@@ -0,0 +1,4 @@
+valkey_host = "127.0.0.1"
+valkey_port = 10002
+
+hf_token = ""

vulntrain/config.py

@@ -0,0 +1,36 @@
+#! /usr/bin/env python
+
+"""This module is responsible for loading the configuration variables."""
+
+import importlib.util
+import os
+
+
+def load_config(path):
+    spec = importlib.util.spec_from_file_location("config", path)
+    if spec:
+        config = importlib.util.module_from_spec(spec)
+        if spec.loader:
+            spec.loader.exec_module(config)
+    return config
+
+
+conf = None
+try:
+    conf = load_config(os.environ.get("VulnTrain_CONFIG", "fedivuln/conf_sample.py"))
+except Exception as exc:
+    raise Exception("No configuration file provided.") from exc
+finally:
+    if not conf:
+        raise Exception("No configuration file provided.")
+
+try:
+    valkey_host = conf.valkey_host
+    valkey_port = conf.valkey_port
+except AttributeError as e:
+    raise Exception(f"Missing configuration variable: {e}")
+
+try:
+    hf_token = conf.hf_token
+except Exception:
+    hf_token = ""

vulntrain/create_dataset.py

@@ -12,6 +12,7 @@
 import valkey
 from datasets import Dataset, DatasetDict  # type: ignore[import-untyped]
 
+from vulntrain.conf import hf_token, valkey_host, valkey_port
 from vulntrain.utils import strip_markdown
 
 
@@ -20,8 +21,8 @@ def __init__(self, sources, nb_rows):
         self.sources = sources
         self.nb_rows = nb_rows
         self.valkey_client = valkey.Valkey(
-            host="127.0.0.1",
-            port=10002,
+            host=valkey_host,
+            port=valkey_port,
             decode_responses=True,
         )
 
@@ -237,7 +238,12 @@ def gen():
 
     print(dataset_dict)
     if args.upload:
-        dataset_dict.push_to_hub(args.repo_id, commit_message=args.commit_message)
+        if hf_token:
+            dataset_dict.push_to_hub(args.repo_id, commit_message=args.commit_message)
+        else:
+            dataset_dict.push_to_hub(
+                args.repo_id, commit_message=args.commit_message, token=hf_token
+            )
 
 
 if __name__ == "__main__":

vulntrain/utils.py

@@ -1,5 +1,5 @@
 from markdown_it import MarkdownIt
-from nltk.tokenize import sent_tokenize
+from nltk.tokenize import sent_tokenize  # type: ignore[import-untyped]
 
 
 def sentences(text, num_sentences=5):