import os from typing import List, Dict import duckdb import pandas as pd from huggingface_hub import hf_hub_download # Initialisations REPO_ID = "Loren/articles_database" cache_dir = "/tmp" os.makedirs(cache_dir, exist_ok=True) # Rediriger le cache HF globalement os.environ["HF_HOME"] = cache_dir os.environ["HF_DATASETS_CACHE"] = cache_dir os.environ["TRANSFORMERS_CACHE"] = cache_dir # Téléchargement des fichiers Parquet depuis Hugging Face articles_parquet = hf_hub_download( repo_id=REPO_ID, filename="articles.parquet", repo_type="dataset", cache_dir=cache_dir) tags_parquet = hf_hub_download( repo_id=REPO_ID, filename="tags.parquet", repo_type="dataset", cache_dir=cache_dir) tag_article_parquet = hf_hub_download( repo_id=REPO_ID, filename="tag_article.parquet", repo_type="dataset", cache_dir=cache_dir) # Connexion DuckDB en mémoire con = duckdb.connect() # Créer des tables DuckDB directement à partir des fichiers Parquet con.execute(f"CREATE VIEW articles AS SELECT * FROM parquet_scan('{articles_parquet}')") con.execute(f"CREATE VIEW tags AS SELECT * FROM parquet_scan('{tags_parquet}')") con.execute(f"CREATE VIEW tag_article AS SELECT * FROM parquet_scan('{tag_article_parquet}')") # Fonctions d'accès aux données def fetch_tags() -> List[str]: """ Récupère la liste de tous les tags disponibles dans la base de données. Returns: List[str]: Une liste de chaînes de caractères correspondant aux noms des tags, triés par ordre alphabétique. """ query = "SELECT tag_name FROM tags ORDER BY tag_name" result = con.execute(query).fetchall() return [row[0] for row in result] def fetch_articles_by_tags(tags: List[str]) -> List[Dict]: """ Récupère les articles associés à un ou plusieurs tags. Args: tags (List[str]): Une liste de noms de tags pour filtrer les articles. Returns: List[Dict]: Une liste de dictionnaires, chacun représentant un article avec les clés: - 'article_id': ID de l'article - 'article_title': Titre de l'article - 'article_url': URL de l'article Notes: - Si la liste `tags` est vide, la fonction retourne une liste vide. - Les résultats incluent uniquement les articles correspondant à au moins un des tags fournis. """ if not tags: return [] placeholders = ",".join(["?"] * len(tags)) query = f"""SELECT distinct a.article_id, a.article_title, a.article_url FROM tags t, tag_article ta, articles a WHERE t.tag_id = ta.tag_id AND ta.article_id = a.article_id AND t.tag_name IN ({placeholders}) """ result = con.execute(query, tags).fetchdf() return result.to_dict(orient="records")