File size: 2,862 Bytes
05ef53c 84fa67b 05ef53c 84fa67b 05ef53c 84fa67b 05ef53c 84fa67b 05ef53c 84fa67b 05ef53c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 | import gradio as gr
from online.graph_rag import GraphRAGApp
# Initialize backend (once)
app = GraphRAGApp(
emb_path="./artifacts/embeddings/embs_epochs200_lr0.0001_weight_decay0.0001_hidden_dim128_num_layer3_neg_k5_margin0.2_lambda_align0.03.pt",
csv_path="./artifacts/graph_triplets/graph_edges.csv",
topk=6,
semantic_contribute=0.5,
score_tolerance_factor=0.4
)
# Streaming inference logic
def answer_query(query):
if not query or query.strip() == "":
yield "Please enter a valid question."
return
# Immediate UI feedback
yield "Thinking...", gr.update(interactive=False)
REASONING_START_TAG = "<think>"
START_TAG = "</think>\n\n<answer>"
# Store stream outputs from LLM with buffer, and start displaying result after the signal </think>\n\n<answer>
buffer = ""
start_yielding = False
has_think_tag = True
stream_output = app.answer(query, stream=True) #Streaming LLM output
for raw_chunk in stream_output:
chunk = raw_chunk["choices"][0]["text"]
buffer += chunk
# Handle when LLM didn't return <think>...</think> block
if has_think_tag:
if len(buffer) > 2 * len(REASONING_START_TAG) and REASONING_START_TAG not in buffer:
start_yielding = True
has_think_tag = False
# Un-comment this for debugging
# print(chunk, end="", flush=True)
if start_yielding:
yield buffer, gr.update(interactive=False)
continue
idx = buffer.find(START_TAG, len(buffer) // 2)
if idx != -1:
buffer = buffer[idx+len(START_TAG):]
start_yielding = True
if len(buffer) == 0:
buffer = "I don't know"
yield buffer, gr.update(interactive=True)
# Gradio UI
with gr.Blocks() as demo:
gr.Markdown("# GNN-based GraphRAG for Healthcare QA")
gr.Markdown(
"Ask a healthcare-related question. "
"The system retrieves knowledge from a graph and generates an answer."
)
with gr.Row():
query_input = gr.Textbox(
label="Your Question",
placeholder="e.g., What is anatomy?"
)
gr.Examples(
examples=[
"What is anatomy ?",
"Is heart a part of cardiovascular system ?",
"What is related to the central nervous system ?"
],
inputs=query_input
)
with gr.Row():
submit_btn = gr.Button("Submit")
with gr.Row():
output = gr.Textbox(
label="Answer",
lines=5
)
# Streaming happens automatically because answer_query uses `yield`
submit_btn.click(
fn=answer_query,
inputs=query_input,
outputs=[output, submit_btn]
)
# Launch
if __name__ == "__main__":
demo.launch() |