Upload create_dataset.py

script/create_dataset.py

@@ -31,20 +31,25 @@ from collections import Counter
 print("Initialisations ...")
 load_dotenv()
 HF_TOKEN = os.getenv('API_HF_TOKEN')
+
+# Constantes
+MIN_COUNT = 5 # nombre minimum d'occurrences pour qu'un tag soit conservé
 DATA_DIR = Path("../../Data")   # dossier parent du script
-REPO_ID = "Loren/articles_db"  # dataset HF
+REPO_ID_DB = "Loren/articles_db"  # dataset HF
+REPO_ID = "Loren/articles_database"  # dataset HF
+DB_NAME = 'articles.db'
+SQLITE_FILE = DATA_DIR / DB_NAME
+LIST_TABLES = ["articles", "tags", "tag_article"]
+PARQUET_DIR = DATA_DIR / "parquet_tables"
 
-parquet_path = hf_hub_download(repo_id=REPO_ID,
+# Chargement des données
+parquet_path = hf_hub_download(repo_id=REPO_ID_DB,
                                filename="medium_articles.parquet",
                                repo_type="dataset")
 
-DB_NAME = 'articles.db'
-SQLITE_FILE = DATA_DIR / DB_NAME
-# Créer le dossier data s'il n'existe pas
+# Créer les dossiers s'ils n'existent pas
 DATA_DIR.mkdir(exist_ok=True)
-list_tables = ["articles", "tags", "tag_article"]
-parquet_dir = DATA_DIR / "parquet_tables"
-parquet_dir.mkdir(exist_ok=True)
+PARQUET_DIR.mkdir(exist_ok=True)
 
 # Chargement des données
 print("Chargement des données ...")
@@ -94,7 +99,7 @@ all_tags = list(itertools.chain.from_iterable(df['list_tags']))
 # Comptage du nombre d'occurrences de chaque tag
 tag_counts = Counter(all_tags)
 # On ne va conserver que les tags avec au moins 100 occurrences
-list_tags = [tag for tag, count in tag_counts.items() if count >= 100]
+list_tags = [tag for tag, count in tag_counts.items() if count >= MIN_COUNT]
 
 # Insertion des tags dans la table
 print("Insertion des tags dans la table ...")
@@ -135,7 +140,7 @@ for _, row in df.iterrows():
             pass
 
 print("-> ", len(list_tags), " tags")
-cur.execute("SELECT COUNT(*) FROM ma_table")
+cur.execute("SELECT COUNT(*) FROM tag_article")
 nb_lignes = cur.fetchone()[0]
 print("-> ", nb_lignes, " associations articles <-> tags")
 print("-> ", len(df), " articles")
@@ -149,7 +154,7 @@ print("Upload base Sqlite dans le dataset hugging face ...")
 upload_file(
     path_or_fileobj=SQLITE_FILE,
     path_in_repo=DB_NAME,
-    repo_id=REPO_ID,
+    repo_id=REPO_ID_DB,
     repo_type="dataset",
     token=HF_TOKEN
 )
@@ -157,9 +162,9 @@ upload_file(
 # Création des fichiers Parquet compressés
 print("Création des fichiers Parquet compressés ...")
 parquet_files = []
-for table in list_tables:
+for table in LIST_TABLES:
     df = pd.read_sql_query(f"SELECT * FROM {table}", conn)
-    parquet_path = parquet_dir / f"{table}.parquet"
+    parquet_path = PARQUET_DIR / f"{table}.parquet"
     df.to_parquet(parquet_path, engine="pyarrow", index=False, compression="snappy")
     parquet_files.append(parquet_path)