Spaces:

Loren
/

api_search_articles

Sleeping

App Files Files Community

Loren commited on 26 days ago

Commit

fe53425

verified ·

1 Parent(s): 1c255df

Update app/main.py

Browse files

Files changed (1) hide show

app/main.py +226 -245

app/main.py CHANGED Viewed

@@ -1,246 +1,227 @@
-from fastapi import FastAPI, Query
-from typing import List, Dict, Any
-from app import database
-from fastapi.middleware.cors import CORSMiddleware
-from fastapi.responses import HTMLResponse
-from pydantic import BaseModel
-from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
-import torch
-from app.templates.prompt_mistral_rag import RAG_PROMPT_TEMPLATE
-app = FastAPI(
-    title="Articles API",
-    description="API pour récupérer articles et tags depuis SQLite",
-    version="1.0"
-)
-# Chargement du modèle génératif
-#MODEL_NAME = "mistralai/Mistral-7B-Instruct-v0.3"
-#tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-#model = AutoModelForCausalLM.from_pretrained(MODEL_NAME,
-#                                             torch_dtype=torch.float16,
-#                                             device_map="auto"
-#                                            )
-model_id = "mistralai/Mistral-7B-Instruct-v0.2"
-# Charger le tokenizer
-tokenizer = AutoTokenizer.from_pretrained(model_id)
-# Config de quantization moderne (4-bit ou 8-bit)
-quant_config = BitsAndBytesConfig(
-    load_in_4bit=True,          # False pour int8
-    bnb_4bit_compute_dtype=torch.float16,  # dtype des calculs
-    bnb_4bit_use_double_quant=True,
-    bnb_4bit_quant_type="nf4"
-)
-# Charger le modèle avec la nouvelle API
-model = AutoModelForCausalLM.from_pretrained(
-    model_id,
-    quantization_config=quant_config,
-    device_map="auto",          # pour GPU auto
-    dtype=torch.float16
-)
-# CORS pour permettre l'accès depuis le navigateur
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],  # autorise toutes les origines
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
-@app.get("/", response_class=HTMLResponse)
-def home():
-    return """
-    <html>
-        <head><title>Page d'accueil</title></head>
-        <body>
-            <h1>Welcome on the API search articles !</h1>
-        </body>
-    </html>
-    """
-@app.get("/get_tags")
-def get_tags():
-    """
-    Récupère la liste de tous les tags disponibles via l'API.
-    Returns:
-        Dict: Un dictionnaire contenant soit la liste des tags, soit les informations d'erreur.
-            - Si succès :
-                {
-                    "status": "ok",
-                    "tags": List[str]  # Liste des noms de tags triés par ordre alphabétique
-                }
-            - En cas d'erreur :
-                {
-                    "status": "error",
-                    "code": str,      # Nom de l'exception
-                    "message": str    # Message de l'exception
-                }
-    Notes:
-        - L'appel de cet endpoint effectue un accès à la base de données via la fonction `fetch_tags`.
-        - En cas de problème avec la base de données, un message d'erreur détaillé est retourné.
-    """
-    try:
-        dict_result = database.fetch_tags()
-        if dict_result["status"] == "ok":
-            return {"status": "ok", "tags": dict_result["result"]}
-        else:
-            return dict_result
-    except Exception as e:
-        return {"status": "error", "code": type(e).__name__, "message": str(e)}
-@app.get("/get_articles_with_tags")
-def get_articles_with_tags(tags: List[str] = Query(..., description="Liste des tags à filtrer")):
-    """
-    Récupère les articles associés à une ou plusieurs tags spécifiés.
-    Args:
-        tags (List[str]): Liste des noms de tags utilisés pour filtrer les articles.
-                          Doit contenir au moins un tag.
-    Returns:
-        Dict: Un dictionnaire contenant soit les articles correspondants, soit les informations d'erreur.
-            - Si succès :
-                {
-                    "status": "ok",
-                    "tags": List[str],        # Tags utilisés pour filtrer
-                    "articles": List[Dict]    # Liste des articles correspondants
-                }
-                Chaque article est un dictionnaire contenant :
-                    - 'article_id': int, ID de l'article
-                    - 'article_title': str, Titre de l'article
-                    - 'article_url': str, URL de l'article
-            - En cas d'erreur :
-                {
-                    "status": "error",
-                    "code": str,             # Code d'erreur ou nom de l'exception
-                    "message": str           # Message d'erreur
-                }
-    Notes:
-        - Si la liste `tags` est vide, la fonction retourne une erreur avec le code "no_tags".
-        - L'appel de cet endpoint utilise la fonction `fetch_articles_by_tags` pour récupérer les articles.
-    """
-    try:
-        dict_result = database.fetch_articles_by_tags(tags)
-        if dict_result["status"] == "ok":
-            return {"status": "ok",
-                    "tags": tags,
-                    "articles": dict_result["result"]}
-        else:
-            return dict_result
-    except Exception as e:
-        return {"status": "error", "code": type(e).__name__, "message": str(e)}
-@app.get("/get_query_results")
-def get_query_results(query: str = Query(..., description="Requête de recherche textuelle"),
-                      k_model: int = Query(10, description="Nombre de candidats retournés par FAISS"),
-                      k_cross: int = Query(5, description="Nombre de résultats conservés après reranking"),
-                      use_rerank: bool = Query(True, description="Indique si le reranking avec cross-encoder doit être utilisé")
-                      ) -> Dict[str, Any]:
-    """
-    Récupère les résultats d'une requête en utilisant deux modèles de recherche.
-    Args:
-        query (str): La requête utilisateur pour laquelle récupérer les résultats.
-        k_model (int, optionel): Nombre de résultats à retourner pour le modèle principal. Par défaut à 10.
-        k_cross (int, optionel): Nombre de résultats à retourner pour le modèle croisé. Par défaut à 5.
-        use_rerank (bool, optionnel): Indique si le reranking avec cross-encoder doit être utilisé. Par défaut à True.
-        Si False, on désactive complètement le cross-encoder et le rerank.
-    Returns:
-        Dict[str, Any]: Un dictionnaire contenant soit les résultats de la requête, soit les informations d'erreur.
-    Notes:
-        - L'appel de cet endpoint utilise la fonction `fetch_query_result` pour obtenir les résultats.
-        - En cas de problème lors du traitement de la requête, un message d'erreur détaillé est retourné.
-    """
-    try:
-        dict_result = database.fetch_query_results(query, k_model, k_cross, use_rerank)
-        if dict_result["status"] == "ok":
-            return {"status": "ok",
-                    "results": dict_result["result"]}
-        else:
-            return dict_result
-    except Exception as e:
-        return {"status": "error", "code": type(e).__name__, "message": str(e)}
-#
-class QueryRequest(BaseModel):
-    question: str
-    use_rerank: bool = True
-@app.post("/get_answer")
-async def ask_question(request: QueryRequest):
-    """
-    Traite une question utilisateur en effectuant une recherche dans la base de données
-    puis en générant une réponse à l’aide du modèle de langage (RAG).
-    Le fonctionnement se déroule en trois étapes :
-    1. Extraction et nettoyage de la question utilisateur.
-    2. Recherche des passages pertinents dans la base de données (`fetch_query_results`).
-    3. Génération d’une réponse fondée sur les morceaux de texte récupérés (RAG).
-    Paramètres
-    ----------
-    request : QueryRequest
-        Objet contenant la question utilisateur sous forme de chaîne de caractères.
-    Retour
-    ------
-    dict
-        Un dictionnaire contenant :
-        - "status" : "ok" si la requête a réussi, "error" en cas d'échec.
-        - "results" : liste des chunks retournés par la base de données (présent seulement si status = "ok").
-        - "answer" : réponse générée par le modèle ou message d’erreur.
-    Exceptions
-    ----------
-    Toute exception survenant durant l’exécution est interceptée
-    et retournée sous forme d’un message d’erreur dans la clé "answer".
-    Notes
-    -----
-    - Si aucun chunk pertinent n'est trouvé, la fonction renvoie un message indiquant
-      que seules les questions relatives aux articles du jeu de données peuvent être traitées.
-    - La génération de la réponse utilise un template RAG et produit jusqu’à 500 tokens.
-    """
-    try:
-        user_query = request.question.strip()
-        use_rerank = request.use_rerank
-        dict_result = database.fetch_query_results(user_query, k_model=10,
-                                                   k_cross=5, use_rerank=use_rerank)
-        if dict_result["status"] == "ok":
-            list_chunks = [resp['chunk_text'] for resp in dict_result['result']]
-            if not list_chunks:
-                answer = ("Je ne dispose pas d’informations sur ce sujet. "
-                          "Je peux uniquement répondre à des questions sur les articles " \
-                          "du jeu de données.")
-            else:
-                # Construction du prompt
-                prompt = RAG_PROMPT_TEMPLATE.format(
-                    context="\n".join(list_chunks),
-                    question=user_query
-                )
-                # Génération de la réponse
-                inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-                outputs = model.generate(**inputs, max_new_tokens=500)
-                generated_tokens = outputs[0][inputs["input_ids"].shape[-1]:]  # uniquement la partie générée
-                answer = tokenizer.decode(generated_tokens, skip_special_tokens=True).strip()
-            return {"status": "ok",
-                    "results": dict_result["result"],
-                    "answer": answer}
-        else:
-            answer = f"Une erreur est survenue lors de la récupération des informations : \
-                       {dict_result['code']} - {dict_result['message']}."
-            return {"status": "error", "answer": answer}
-    except Exception as e:
-        answer = f"Une erreur est survenue lors de la récupération des informations : \
-                       {type(e).__name__} - {str(e)}."
         return {"status": "error", "answer": answer}

+from fastapi import FastAPI, Query
+from typing import List, Dict, Any
+from app import database
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import HTMLResponse
+from pydantic import BaseModel
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
+from app.templates.prompt_mistral_rag import RAG_PROMPT_TEMPLATE
+app = FastAPI(
+    title="Articles API",
+    description="API pour récupérer articles et tags depuis SQLite",
+    version="1.0"
+)
+# Chargement du modèle génératif
+MODEL_NAME = "mistralai/Mistral-7B-Instruct-v0.3"
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+model = AutoModelForCausalLM.from_pretrained(MODEL_NAME,
+                                             torch_dtype=torch.float16,
+                                             device_map="auto"
+                                            )
+# CORS pour permettre l'accès depuis le navigateur
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],  # autorise toutes les origines
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+@app.get("/", response_class=HTMLResponse)
+def home():
+    return """
+    <html>
+        <head><title>Page d'accueil</title></head>
+        <body>
+            <h1>Welcome on the API search articles !</h1>
+        </body>
+    </html>
+    """
+@app.get("/get_tags")
+def get_tags():
+    """
+    Récupère la liste de tous les tags disponibles via l'API.
+    Returns:
+        Dict: Un dictionnaire contenant soit la liste des tags, soit les informations d'erreur.
+            - Si succès :
+                {
+                    "status": "ok",
+                    "tags": List[str]  # Liste des noms de tags triés par ordre alphabétique
+                }
+            - En cas d'erreur :
+                {
+                    "status": "error",
+                    "code": str,      # Nom de l'exception
+                    "message": str    # Message de l'exception
+                }
+    Notes:
+        - L'appel de cet endpoint effectue un accès à la base de données via la fonction `fetch_tags`.
+        - En cas de problème avec la base de données, un message d'erreur détaillé est retourné.
+    """
+    try:
+        dict_result = database.fetch_tags()
+        if dict_result["status"] == "ok":
+            return {"status": "ok", "tags": dict_result["result"]}
+        else:
+            return dict_result
+    except Exception as e:
+        return {"status": "error", "code": type(e).__name__, "message": str(e)}
+@app.get("/get_articles_with_tags")
+def get_articles_with_tags(tags: List[str] = Query(..., description="Liste des tags à filtrer")):
+    """
+    Récupère les articles associés à une ou plusieurs tags spécifiés.
+    Args:
+        tags (List[str]): Liste des noms de tags utilisés pour filtrer les articles.
+                          Doit contenir au moins un tag.
+    Returns:
+        Dict: Un dictionnaire contenant soit les articles correspondants, soit les informations d'erreur.
+            - Si succès :
+                {
+                    "status": "ok",
+                    "tags": List[str],        # Tags utilisés pour filtrer
+                    "articles": List[Dict]    # Liste des articles correspondants
+                }
+                Chaque article est un dictionnaire contenant :
+                    - 'article_id': int, ID de l'article
+                    - 'article_title': str, Titre de l'article
+                    - 'article_url': str, URL de l'article
+            - En cas d'erreur :
+                {
+                    "status": "error",
+                    "code": str,             # Code d'erreur ou nom de l'exception
+                    "message": str           # Message d'erreur
+                }
+    Notes:
+        - Si la liste `tags` est vide, la fonction retourne une erreur avec le code "no_tags".
+        - L'appel de cet endpoint utilise la fonction `fetch_articles_by_tags` pour récupérer les articles.
+    """
+    try:
+        dict_result = database.fetch_articles_by_tags(tags)
+        if dict_result["status"] == "ok":
+            return {"status": "ok",
+                    "tags": tags,
+                    "articles": dict_result["result"]}
+        else:
+            return dict_result
+    except Exception as e:
+        return {"status": "error", "code": type(e).__name__, "message": str(e)}
+@app.get("/get_query_results")
+def get_query_results(query: str = Query(..., description="Requête de recherche textuelle"),
+                      k_model: int = Query(10, description="Nombre de candidats retournés par FAISS"),
+                      k_cross: int = Query(5, description="Nombre de résultats conservés après reranking"),
+                      use_rerank: bool = Query(True, description="Indique si le reranking avec cross-encoder doit être utilisé")
+                      ) -> Dict[str, Any]:
+    """
+    Récupère les résultats d'une requête en utilisant deux modèles de recherche.
+    Args:
+        query (str): La requête utilisateur pour laquelle récupérer les résultats.
+        k_model (int, optionel): Nombre de résultats à retourner pour le modèle principal. Par défaut à 10.
+        k_cross (int, optionel): Nombre de résultats à retourner pour le modèle croisé. Par défaut à 5.
+        use_rerank (bool, optionnel): Indique si le reranking avec cross-encoder doit être utilisé. Par défaut à True.
+        Si False, on désactive complètement le cross-encoder et le rerank.
+    Returns:
+        Dict[str, Any]: Un dictionnaire contenant soit les résultats de la requête, soit les informations d'erreur.
+    Notes:
+        - L'appel de cet endpoint utilise la fonction `fetch_query_result` pour obtenir les résultats.
+        - En cas de problème lors du traitement de la requête, un message d'erreur détaillé est retourné.
+    """
+    try:
+        dict_result = database.fetch_query_results(query, k_model, k_cross, use_rerank)
+        if dict_result["status"] == "ok":
+            return {"status": "ok",
+                    "results": dict_result["result"]}
+        else:
+            return dict_result
+    except Exception as e:
+        return {"status": "error", "code": type(e).__name__, "message": str(e)}
+#
+class QueryRequest(BaseModel):
+    question: str
+    use_rerank: bool = True
+@app.post("/get_answer")
+async def ask_question(request: QueryRequest):
+    """
+    Traite une question utilisateur en effectuant une recherche dans la base de données
+    puis en générant une réponse à l’aide du modèle de langage (RAG).
+    Le fonctionnement se déroule en trois étapes :
+    1. Extraction et nettoyage de la question utilisateur.
+    2. Recherche des passages pertinents dans la base de données (`fetch_query_results`).
+    3. Génération d’une réponse fondée sur les morceaux de texte récupérés (RAG).
+    Paramètres
+    ----------
+    request : QueryRequest
+        Objet contenant la question utilisateur sous forme de chaîne de caractères.
+    Retour
+    ------
+    dict
+        Un dictionnaire contenant :
+        - "status" : "ok" si la requête a réussi, "error" en cas d'échec.
+        - "results" : liste des chunks retournés par la base de données (présent seulement si status = "ok").
+        - "answer" : réponse générée par le modèle ou message d’erreur.
+    Exceptions
+    ----------
+    Toute exception survenant durant l’exécution est interceptée
+    et retournée sous forme d’un message d’erreur dans la clé "answer".
+    Notes
+    -----
+    - Si aucun chunk pertinent n'est trouvé, la fonction renvoie un message indiquant
+      que seules les questions relatives aux articles du jeu de données peuvent être traitées.
+    - La génération de la réponse utilise un template RAG et produit jusqu’à 500 tokens.
+    """
+    try:
+        user_query = request.question.strip()
+        use_rerank = request.use_rerank
+        dict_result = database.fetch_query_results(user_query, k_model=10,
+                                                   k_cross=5, use_rerank=use_rerank)
+        if dict_result["status"] == "ok":
+            list_chunks = [resp['chunk_text'] for resp in dict_result['result']]
+            if not list_chunks:
+                answer = ("Je ne dispose pas d’informations sur ce sujet. "
+                          "Je peux uniquement répondre à des questions sur les articles " \
+                          "du jeu de données.")
+            else:
+                # Construction du prompt
+                prompt = RAG_PROMPT_TEMPLATE.format(
+                    context="\n".join(list_chunks),
+                    question=user_query
+                )
+                # Génération de la réponse
+                inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+                outputs = model.generate(**inputs, max_new_tokens=500)
+                generated_tokens = outputs[0][inputs["input_ids"].shape[-1]:]  # uniquement la partie générée
+                answer = tokenizer.decode(generated_tokens, skip_special_tokens=True).strip()
+            return {"status": "ok",
+                    "results": dict_result["result"],
+                    "answer": answer}
+        else:
+            answer = f"Une erreur est survenue lors de la récupération des informations : \
+                       {dict_result['code']} - {dict_result['message']}."
+            return {"status": "error", "answer": answer}
+    except Exception as e:
+        answer = f"Une erreur est survenue lors de la récupération des informations : \
+                       {type(e).__name__} - {str(e)}."
         return {"status": "error", "answer": answer}