Spaces:
Running on Zero
Running on Zero
File size: 5,049 Bytes
b1aba72 db4e733 b1aba72 51e9502 b1aba72 aa74750 b1aba72 | 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 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 | """
app.py
--------------
Agentic RAG — HuggingFace Space entry point (ZeroGPU version).
"""
import spaces # noqa: F401 - MUST stay first; see docstring above
import time
import gradio as gr
from config import CONFIDENCE_THRESHOLD, MAX_PROMPT_TOKENS
from retrieval import run_pipeline_phase1
from generation import generate_answer
from utils import extract_sources, clean_final_answer
from ui_helpers import format_metrics
# Chat handler
def chat(message: str, history: list):
"""
Sync generator for gr.ChatInterface(fn=chat, additional_outputs=[...]).
This function has no GPU-specific code. It calls run_pipeline_phase1() (FAISS +
Tavily, both CPU-only) and generate_answer() (which internally
calls the @spaces.GPU-decorated _generate_on_gpu() exactly once).
Yields 2-tuples: (bot_message_string, metrics_markdown_string).
"""
t_total_start = time.perf_counter()
# Step 1: Immediate placeholder
yield "🔍 Searching...", format_metrics()
# Step 2: Retrieval (FAISS + optional Tavily, CPU-only)
state = run_pipeline_phase1(message, history=history)
t_retrieve = state.get("t_retrieve", 0.0)
t_web = state.get("t_web", 0.0)
confidence = state.get("confidence", 0.0)
source_type = state.get("source_type", "none")
sources_count = state.get("sources_count", 0)
if state.get("needs_web") and sources_count > 0:
status = (
f"🌐 {sources_count} sources fetched ({t_web:.1f}s). "
f"Generating on GPU…"
)
elif state.get("needs_web"):
status = "🌐 Web search done. Generating on GPU…"
else:
status = f"📚 FAISS (sim={confidence:.3f}). Generating on GPU…"
yield status, format_metrics(
source_type = source_type,
confidence = confidence,
sources_count = sources_count,
generating = True,
elapsed = time.perf_counter() - t_total_start,
t_retrieve = t_retrieve,
t_web = t_web,
)
# Step 3: Word-by-word generation
# generate_answer() internally calls the @spaces.GPU-decorated
# _generate_on_gpu() once (blocking - but on ZeroGPU this should
# take a few seconds), then yields the response word by word.
answer = ""
output_tokens = 0
prompt_tokens = 0
t_gen_start = time.perf_counter()
last_metrics = format_metrics(
source_type = source_type,
confidence = confidence,
sources_count = sources_count,
generating = True,
elapsed = 0.0,
t_retrieve = t_retrieve,
t_web = t_web,
)
try:
for word_text, prompt_tok_count in generate_answer(state):
answer += word_text
output_tokens += 1
prompt_tokens = prompt_tok_count
if output_tokens % 5 == 0:
last_metrics = format_metrics(
source_type = source_type,
confidence = confidence,
sources_count = sources_count,
generating = True,
elapsed = time.perf_counter() - t_gen_start,
t_retrieve = t_retrieve,
t_web = t_web,
token_so_far = output_tokens,
)
yield answer, last_metrics
except Exception as e:
error_msg = f"{type(e).__name__}: {e}"
yield f"❌ Generation failed: {error_msg}", format_metrics(error=error_msg)
return
# Step 4: Final answer + sources + metrics
t_generate = time.perf_counter() - t_gen_start
t_total = time.perf_counter() - t_total_start
tokens_per_sec = output_tokens / t_generate if t_generate > 0 else 0.0
sources_section = extract_sources(state)
final_answer = clean_final_answer(answer) + sources_section
yield final_answer, format_metrics(
source_type = source_type,
confidence = confidence,
sources_count = sources_count,
prompt_tokens = prompt_tokens,
t_retrieve = t_retrieve,
t_web = t_web,
t_generate = t_generate,
t_total = t_total,
tokens_per_sec = tokens_per_sec,
output_tokens = output_tokens,
)
# UI Layout
with gr.Blocks(theme=gr.themes.Soft(), title="Agentic RAG (ZeroGPU)") as demo:
gr.Markdown(
"## 🤖 Agentic RAG "
)
gr.Markdown(
"💡**Tip:** Include full context in every question for best results. \n"
"Follow-up questions are supported but specific questions always "
"perform better.\n\n"
"⚡**Note:** Running on a shared free GPU (ZeroGPU). Most responses "
"complete in a few seconds, though you may occasionally wait in a "
"short queue if the shared GPU is busy with other Spaces."
)
with gr.Row():
with gr.Column(scale=3):
metrics_panel = gr.Markdown(value=format_metrics(), render=False)
gr.ChatInterface(
fn=chat,
additional_outputs=[metrics_panel],
)
with gr.Column(scale=1, min_width=220):
gr.Markdown("**📊 Metrics**")
metrics_panel.render()
# Launch
# No share=True — HF Spaces provides its own public URL automatically.
demo.queue(default_concurrency_limit=1)
demo.launch()
|