Update app.py

app.py

@@ -1,44 +1,95 @@
+import json
 import os
-if not os.path.exists("data/chunks.json"):
-    from datasets import load_dataset
-    import json
-    from pathlib import Path
-    Path("data").mkdir(exist_ok=True)
-    print("Downloading dataset...")
-    fb = load_dataset("PatronusAI/financebench", split="train")
-    fa = load_dataset("gbharti/finance-alpaca", split="train")
-    chunks, entities, metadata = [], [], []
-    for i, row in enumerate(fb):
-        for j, ev in enumerate(row["evidence"]):
-            text = ev.get("evidence_text","").strip()
-            if not text: continue
-            chunk_id = f"fb_{i:04d}_{j:02d}"
-            chunks.append({"chunk_id":chunk_id,"company":row["company"],"doc_name":row["doc_name"],"period":row["doc_period"],"text":text,"question":row["question"],"answer":row["answer"],"source":"financebench"})
-            entities.append({"chunk_id":chunk_id,"company":row["company"],"entities":[]})
-            metadata.append({"chunk_id":chunk_id,"company":row["company"],"doc_name":row["doc_name"],"period":row["doc_period"],"question":row["question"],"answer":row["answer"],"source":"financebench"})
-    def chunk_text(text, size=512, overlap=50):
-        words = text.split()
-        results = []
-        start = 0
-        while start < len(words):
-            end = min(start+size, len(words))
-            results.append(" ".join(words[start:end]))
-            if end == len(words): break
-            start += size-overlap
-        return results
-    target = 1_000_000
-    current = sum(len(c["text"].split()) for c in chunks)
-    for i, row in enumerate(fa):
-        if current >= target: break
-        text = row["output"].strip()
-        if len(text.split()) < 20: continue
-        for j, ct in enumerate(chunk_text(text)):
-            chunk_id = f"fa_{i:05d}_{j:02d}"
-            chunks.append({"chunk_id":chunk_id,"company":"general","doc_name":f"alpaca_{i}","period":"general","text":ct,"question":row["instruction"],"answer":row["output"],"source":"alpaca"})
-            entities.append({"chunk_id":chunk_id,"company":"general","entities":[]})
-            metadata.append({"chunk_id":chunk_id,"company":"general","doc_name":f"alpaca_{i}","period":"general","question":row["instruction"],"answer":row["output"],"source":"alpaca"})
-            current += len(ct.split())
-    json.dump(chunks, open("data/chunks.json","w",encoding="utf-8"), indent=2)
-    json.dump(entities, open("data/entities.json","w",encoding="utf-8"), indent=2)
-    json.dump(metadata, open("data/metadata.json","w",encoding="utf-8"), indent=2)
-    print(f"Done. {len(chunks)} chunks")
+from fastapi import FastAPI
+import gradio as gr
+
+# ----------------------------
+# FastAPI App
+# ----------------------------
+app = FastAPI()
+
+@app.get("/")
+def home():
+    return {"status": "GraphRAG running"}
+
+# ----------------------------
+# Load Data
+# ----------------------------
+DATA_DIR = "data"
+
+chunks = []
+entities = []
+metadata = []
+
+try:
+    with open(os.path.join(DATA_DIR, "chunks.json"), "r", encoding="utf-8") as f:
+        chunks = json.load(f)
+
+    with open(os.path.join(DATA_DIR, "entities.json"), "r", encoding="utf-8") as f:
+        entities = json.load(f)
+
+    with open(os.path.join(DATA_DIR, "metadata.json"), "r", encoding="utf-8") as f:
+        metadata = json.load(f)
+
+    print(f"Loaded {len(chunks)} chunks")
+
+except Exception as e:
+    print(f"Error loading data: {e}")
+
+
+# ----------------------------
+# Simple Search Function
+# ----------------------------
+def search_query(query):
+    if not query.strip():
+        return "Enter a query"
+
+    results = []
+
+    query = query.lower()
+
+    for chunk in chunks[:500]:  # limit for speed
+        text = chunk.get("text", "").lower()
+
+        if query in text:
+            results.append(
+                f"Company: {chunk.get('company', 'N/A')}\n"
+                f"Source: {chunk.get('source', 'N/A')}\n"
+                f"Text: {chunk.get('text', '')[:500]}"
+            )
+
+        if len(results) >= 5:
+            break
+
+    if not results:
+        return "No relevant results found."
+
+    return "\n\n" + ("-" * 60 + "\n\n").join(results)
+
+
+# ----------------------------
+# Gradio UI
+# ----------------------------
+with gr.Blocks() as demo:
+    gr.Markdown("# GraphRAG Search")
+
+    query = gr.Textbox(
+        label="Ask a question",
+        placeholder="Enter finance-related query..."
+    )
+
+    output = gr.Textbox(
+        label="Results",
+        lines=20
+    )
+
+    btn = gr.Button("Search")
+    btn.click(search_query, inputs=query, outputs=output)
+
+# ----------------------------
+# Launch
+# ----------------------------
+demo.launch(
+    server_name="0.0.0.0",
+    server_port=7860
+)