Upload app.py

app.py

@@ -24,10 +24,10 @@ def load_articles_list():
     try:
         with open("metadata.pkl", "rb") as f:
             df = pickle.load(f)
-        articles = sorted(df['article_category'].unique().tolist())
-        return ["All categories"] + articles
+        articles = sorted(df['article_title'].unique().tolist())
+        return ["All articles"] + articles
     except:
-        return ["All categories"]
+        return ["All articles"]
 
 ARTICLE_LIST = load_articles_list()
 
@@ -101,10 +101,10 @@ def generate_response(message, history, rag_mode, article_filter, max_tokens, te
         vs = get_vectorstore(backend)
         if vs:
             try:
-                filter_dict = {"article_category": article_filter} if article_filter != "All categories" else None
+                filter_dict = {"article_title": article_filter} if article_filter != "All articles" else None
                 docs = vs.similarity_search(message, k=5, filter=filter_dict)
                 context = "\n\n".join([
-                    f"[Category: {doc.metadata.get('article_category', 'N/A')}] {doc.page_content[:700]}"
+                    f"[Category: {doc.metadata.get('article_title', 'N/A')}] {doc.page_content[:700]}"
                     for doc in docs
                 ])
             except:
@@ -144,9 +144,14 @@ with gr.Blocks(title="Article Q&A model") as demo:
     with gr.Row():
         rag_mode = gr.Radio(
             choices=list(BACKENDS.keys()),
-            value="Fast Mode (No RAG)",
+            value="FAISS - RAG (HNSW)",
             label="Mode"
         )
+        article_filter = gr.Dropdown(
+            choices=ARTICLE_LIST,
+            value="All articles",
+            label="Focus on article"
+        )
 
     with gr.Accordion("Advanced Generation Parameters", open=False):
         max_tokens = gr.Slider(256, 2048, value=900, step=64, label="Max Tokens")