feat: add citations showup

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+vector_store/index.faiss filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,17 @@
+# Python-generated files
+__pycache__/
+*.py[oc]
+build/
+dist/
+wheels/
+*.egg-info
+
+# Env
+.venv
+.env
+
+# UV lock
+uv.lock
+
+# Store
+/vector_store

src/core/config.py

@@ -21,8 +21,9 @@ class Settings(BaseSettings):
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
         # Try to get token from environment if not set
+        # HuggingFace Spaces uses HF_TOKEN by default
         if not self.HUGGINGFACE_TOKEN:
-            self.HUGGINGFACE_TOKEN = os.getenv("HUGGINGFACE_TOKEN", "")
+            self.HUGGINGFACE_TOKEN = os.getenv("HF_TOKEN", os.getenv("HUGGINGFACE_TOKEN", ""))
 
 
 settings = Settings()

src/gradio_app.py

@@ -2,19 +2,52 @@ import gradio as gr
 from src.rag.pipeline import answer_question
 
 
+def format_answer(question: str) -> tuple[str, str]:
+    """Format answer with citations."""
+    answer, citations = answer_question(question)
+    
+    # Format citations
+    if citations:
+        citations_text = "\n\n### 📚 Sources:\n"
+        for i, citation in enumerate(citations, 1):
+            citations_text += f"{i}. **{citation['title']}** ({citation['year']})\n"
+    else:
+        citations_text = ""
+    
+    return answer, citations_text
+
+
 def run_gradio():
-    with gr.Blocks(title="PaperMate") as demo:
-        gr.Markdown("# 📄 PaperMate — Ask about research papers")
-        question = gr.Textbox(
-            label="Enter your question",
-            placeholder="e.g. What is the NEMO paper about?",
-            lines=15,
-            max_lines=30,
+    with gr.Blocks(title="PaperMate", theme=gr.themes.Soft()) as demo:
+        gr.Markdown(
+            """
+            # 📄 PaperMate — Research Paper Q&A Assistant
+            Ask questions about research papers and get answers backed by scientific literature.
+            """
         )
-        output = gr.Textbox(label="Answer", lines=15, max_lines=30)
+        
+        with gr.Row():
+            with gr.Column():
+                question = gr.Textbox(
+                    label="Your Question",
+                    placeholder="e.g., What techniques are used to handle out-of-vocabulary words in NLP?",
+                    lines=3,
+                )
+                btn = gr.Button("🔍 Search & Answer", variant="primary")
+        
+        with gr.Row():
+            with gr.Column():
+                output = gr.Textbox(label="Answer", lines=10, max_lines=20)
+                citations = gr.Markdown(label="Sources")
 
-        btn = gr.Button("Ask")
-        btn.click(fn=answer_question, inputs=question, outputs=output)
+        btn.click(fn=format_answer, inputs=question, outputs=[output, citations])
+        
+        gr.Markdown(
+            """
+            ---
+            💡 **Tip:** Questions are answered using relevant papers from the ArXiv dataset.
+            """
+        )
 
     demo.launch()
 

src/prompts.py

@@ -1,10 +1,18 @@
 SYSTEM_PROMPT = """
-You are a helpful assistant that provides accurate and concise information about scientific papers based on the given context.
-Do not provide any information that is not included in the context. Do not mention context details in your answer.
+You are a knowledgeable research assistant that provides accurate, well-structured answers about scientific papers.
+
+Guidelines:
+- Base your answer ONLY on the information provided in the context below
+- Provide clear, concise, and informative responses
+- When relevant, mention specific findings or methodologies from the papers
+- Do NOT make up information or reference papers not in the context
+- Do NOT mention "the context" or "the provided papers" in your response - answer naturally
 
 Context:
 {context}
 
 User Question:
 {question}
+
+Provide a comprehensive answer based on the information above:
 """

src/rag/pipeline.py

@@ -21,8 +21,23 @@ def create_context(docs) -> str:
     return context
 
 
-def answer_question(question: str) -> str:
-    """Answer a question using retrieved and reranked documents."""
+def extract_citations(docs) -> list[dict]:
+    """Extract paper titles and years for citations."""
+    citations = []
+    for doc in docs:
+        citations.append({
+            "title": doc.metadata.get('Titles', 'No Title'),
+            "year": doc.metadata.get('Years', 'Unknown')
+        })
+    return citations
+
+
+def answer_question(question: str) -> tuple[str, list[dict]]:
+    """Answer a question using retrieved and reranked documents.
+    
+    Returns:
+        tuple: (answer, citations) where citations is a list of dicts with 'title' and 'year'
+    """
     # Retrieve more documents than needed for reranking
     try:
         retriever.k = settings.RETRIEVER_K_BEFORE_RERANK
@@ -34,12 +49,13 @@ def answer_question(question: str) -> str:
         )
 
         context = create_context(reranked_results)
+        citations = extract_citations(reranked_results)
         logging.info(f"Constructed context for LLM: {context}")
         chain = get_chain()
 
         response = chain.invoke({"context": context, "question": question})
+        return response, citations
     except Exception as e:
         logging.error(f"Error occurred while answering question: {e}")
-        response = "Sorry, exception occurred while processing your request. See logs for details."
-
-    return response
+        error_msg = f"Sorry, an error occurred while processing your request: {str(e)}"
+        return error_msg, []