api_search_articles / script /create_sqlite_db.py
Loren's picture
Upload 2 files
e8ff22d verified
raw
history blame
3.85 kB
import sqlite3
import pandas as pd
import os
import itertools
import ast
import uuid
from huggingface_hub import hf_hub_download, upload_file
from pathlib import Path
# Initialisations
print("Initialisations ...")
hf_token = os.environ["API_HF_TOKEN"]
DATA_DIR = Path("../data") # dossier parent du script
parquet_path = hf_hub_download(repo_id="Loren/articles_data",
filename="mon_fichier.parquet")
REPO_ID = "Loren/articles_db" # dataset HF
DB_NAME = 'articles.db'
SQLITE_FILE = DATA_DIR / DB_NAME
# Créer le dossier data s'il n'existe pas
DATA_DIR.mkdir(exist_ok=True)
# Chargement des données
print("Chargement des données ...")
df = pd.read_parquet(parquet_path)
# Initialisations de la base SQLite
print("Initialisations de la base SQLite ...")
conn = sqlite3.connect(SQLITE_FILE)
cur = conn.cursor()
# Suppression des anciennes tables
cur.execute("DROP TABLE IF EXISTS tag_article")
cur.execute("DROP TABLE IF EXISTS tags")
cur.execute("DROP TABLE IF EXISTS articles")
# Création des tables Articles, Tags, et de la table d'association articles <-> tags
cur.execute("""
CREATE TABLE articles (
article_id TEXT PRIMARY KEY, -- UUID
article_title TEXT,
article_text TEXT,
article_url TEXT,
article_authors TEXT,
article_date TEXT -- YYYY-MM-DD
)""")
cur.execute("""
CREATE TABLE tags (
tag_id INTEGER PRIMARY KEY AUTOINCREMENT,
tag_name TEXT UNIQUE
)""")
cur.execute("""
CREATE TABLE tag_article (
tag_article_id INTEGER PRIMARY KEY AUTOINCREMENT,
article_id TEXT,
tag_id INTEGER,
FOREIGN KEY(article_id) REFERENCES articles(article_id),
FOREIGN KEY(tag_id) REFERENCES tags(tag_id)
)""")
# Extraction des tags en une liste unique
print("Extraction des tags en une liste unique ...")
df['list_tags'] = df['tags'].apply(lambda x: ast.literal_eval(x) if isinstance(x, str) else [])
# Extraire tous les tags uniques
all_tags = list(set(itertools.chain.from_iterable(df['list_tags'])))
# Insertion des tags dans la table
print("Insertion des tags dans la table ...")
cur.executemany("INSERT INTO tags (tag_name) VALUES (?)", [(tag,) for tag in all_tags])
# Récupération des correspondances tag_name -> tag_id
print("Récupération des correspondances tag_name -> tag_id ...")
cur.execute("SELECT tag_id, tag_name FROM tags")
dict_tag_map = {tag_name: tag_id for tag_id, tag_name in cur.fetchall()}
# Insertion des articles et table d'association dans les tables
print("Insertion des articles et table d'association dans les tables ...")
for _, row in df.iterrows():
# Détermination de l'id article
article_id = str(uuid.uuid4())
# Extraction de la date du timestamp
date_value = None
if pd.notna(row["timestamp"]):
try:
date_value = str(pd.to_datetime(row["timestamp"]).date())
except Exception:
date_value = None
# Insertion dans la table Articles
cur.execute("""
INSERT INTO articles (article_id, article_title, article_text, article_url, article_authors, article_date)
VALUES (?, ?, ?, ?, ?, ?)""",
(article_id, row["title"], row["text"], row["url"], row["authors"], date_value))
# Association aux tags
for tag_name in row['list_tags']:
tag_id = dict_tag_map[tag_name]
cur.execute("INSERT INTO tag_article (article_id, tag_id) VALUES (?, ?)",
(article_id, tag_id))
conn.commit()
conn.close()
# Upload dans le dataset hugging face
print("Upload dans le dataset hugging face ...")
upload_file(
path_or_fileobj=SQLITE_FILE,
path_in_repo=DB_NAME,
repo_id=REPO_ID,
repo_type="dataset",
token=hf_token
)
print("Traitement terminé.")