"""
app.py — Gradio entry point for the Feature Finder Space.
Flow per request:
1. Parse repo URL, check cache (24h TTL)
2. CPU: fetch GitHub commits + Gemini summarization → team_context
3. CPU: embedding pre-filter against precomputed arXiv pool → top-50
4. GPU (@spaces.GPU): v1.4 scoring of top-50, then base-model PR body
5. Render: paper card, confidence tier, PR body, alternates, CTA
The @spaces.GPU function is intentionally narrow: just scoring +
generation. Everything else runs on CPU, keeping Zero-GPU quota usage
proportional to actual model work.
"""
from __future__ import annotations
import os
import time
import traceback
from typing import Any
import gradio as gr
import numpy as np
import cache
import candidate_filter
import confidence
import model_io
import team_context
# The @spaces.GPU decorator is for Zero-GPU Spaces only — it spawns a
# worker process that requests a GPU from the Zero-GPU pool. On
# persistent paid GPU Spaces (h200, a10g-small, etc.), the GPU is
# always attached to the main container; the decorator would either be
# a no-op or fail at worker-init.
#
# Set USE_SPACES_GPU=false on persistent GPU Spaces. Default true so
# Zero-GPU deployments keep working unchanged.
USE_SPACES_GPU = os.environ.get("USE_SPACES_GPU", "true").lower() in ("true", "1", "yes")
if USE_SPACES_GPU:
import spaces
_gpu_decorator = spaces.GPU(duration=60)
else:
_gpu_decorator = lambda f: f # no-op on persistent GPU
# Loaded at cold-start (CPU). The GPU model is loaded inside the
# scoring function on first invocation.
EMBEDDER_PRELOAD = candidate_filter.load_embedder()
POOL_PRELOAD = candidate_filter.load_pool()
@_gpu_decorator
def gpu_score_and_maybe_generate(team_ctx: str, candidates_records: list[dict],
do_local_generation: bool) -> dict:
"""The GPU-only step. Always scores candidates with v1.4 LoRA enabled.
Conditionally generates the PR body locally (LoRA disabled) if
do_local_generation=True. When the configured backend is Gemini,
pass False — the orchestrator will call Gemini on CPU after this
function returns, freeing GPU time for the next request.
Each candidate record is expected to include an `embedding_similarity`
field (the cosine to the team's domain summary), carried through from
the CPU pre-filter. confidence.tier_score uses it as an OOD detector."""
model, tokenizer = model_io.load()
# Score every candidate with LoRA active
model.enable_adapters()
scores = []
for c in candidates_records:
paper_text = f"Title: {c['title']}\n\nAbstract: {c['abstract']}"
s = model_io.score_pair(team_ctx, paper_text)
scores.append(s)
scores_arr = np.array(scores, dtype=np.float32)
top_idx = int(np.argmax(scores_arr))
top = candidates_records[top_idx]
top_score = float(scores_arr[top_idx])
top_sim = float(top.get("embedding_similarity", 0.0))
# Pick alternates: 2nd through 4th by score
order = np.argsort(-scores_arr)
alternates = [candidates_records[i] for i in order[1:4]]
alternate_scores = [float(scores_arr[i]) for i in order[1:4]]
alternate_sims = [float(candidates_records[i].get("embedding_similarity", 0.0)) for i in order[1:4]]
# Tier the top pick AND each alternate against the same pool
# distribution. Pass embedding similarity so OOD vocabulary collisions
# (high z but low cosine) are correctly demoted.
tier = confidence.tier_score(top_score, scores_arr, similarity=top_sim)
alternate_tiers = [
confidence.tier_score(s, scores_arr, similarity=sim)
for s, sim in zip(alternate_scores, alternate_sims)
]
pr_body: str | None = None
if do_local_generation:
# Generate PR body with LoRA off (the base model handles generation)
model.disable_adapters()
paper_text = f"Title: {top['title']}\n\nAbstract: {top['abstract']}"
pr_body = model_io._generate_pr_body_local(team_ctx, paper_text)
model.enable_adapters() # leave in scoring state for next request
return {
"top": top,
"top_score": top_score,
"tier": tier.__dict__,
"alternates": [
{"paper": a, "score": s, "tier": t.__dict__}
for a, s, t in zip(alternates, alternate_scores, alternate_tiers)
],
"pr_body": pr_body, # None when generation is delegated to Gemini
"scores_summary": {
"n_candidates": len(scores),
"max": top_score,
"median": float(np.median(scores_arr)),
"min": float(np.min(scores_arr)),
},
}
def _format_result(result: dict, team_ctx: str, elapsed_s: float) -> tuple[str, str, str, str, str]:
"""Render the per-request result into the 5 Gradio output components."""
top = result["top"]
tier = result["tier"]
arxiv_id = top.get("arxiv_id", "")
url = f"https://arxiv.org/abs/{arxiv_id}" if arxiv_id else ""
tier_emoji = {"high": "🟢", "moderate": "🟡", "low": "🟠", "noise": "🔴"}
top_emoji = tier_emoji.get(tier["tier"], "⚪️")
paper_md = (
f"### {top_emoji} {top['title']}\n\n"
f"[arxiv.org/abs/{arxiv_id}]({url}) · "
f"score {result['top_score']:+.3f} · "
f"margin {tier['margin_z']:+.2f}σ above candidate-pool median\n\n"
f"**Abstract**\n\n{top['abstract']}"
)
confidence_md = (
f"## {top_emoji} {tier['label']}\n\n"
f"**Estimated reliability**: ~{tier['confidence_pct']:.0f}% accurate at this confidence tier\n\n"
f"{tier['explainer']}"
)
# Button-styled CTA link with explicit colors so it renders
# readably in both light and dark Gradio themes (the default
# markdown link blue clashes with dark-mode backgrounds).
cta_button = (
''
'Sign up at engine.remyx.ai →'
)
if tier["tier"] == "high":
# Light footer CTA — user is getting real value from the free
# generalist; don't push hard, but keep the upsell visible.
confidence_md += (
"\n\n---\n\n"
"_Want recommendations tailored to your team's development practices?_\n\n"
f"{cta_button}"
)
else:
# Prominent CTA — user saw the generalist's limits (moderate /
# low / noise tier). Personalized scoring directly addresses the
# cause: a model adapted to the team's development practices.
confidence_md += (
"\n\n---\n\n"
"### Get sharper recommendations on this repo\n\n"
"A personalized model tailored to your team's development practices "
"typically lifts confidence by **10-15pts** on team-specific picks — "
"and produces PR drafts that match your team's actual engineering style.\n\n"
f"{cta_button}"
)
pr_body_md = result["pr_body"]
alternates_md = "### Other candidates\n\n"
for entry in result["alternates"]:
p = entry["paper"]
s = entry["score"]
t = entry["tier"]
emoji = tier_emoji.get(t["tier"], "⚪️")
aid = p.get("arxiv_id", "")
alternates_md += (
f"- {emoji} **{p['title']}** "
f"[arxiv]({f'https://arxiv.org/abs/{aid}' if aid else '#'}) · "
f"{t['label']} · score {s:+.3f} ({t['margin_z']:+.2f}σ)\n"
)
domain_summary = result.get("domain_summary", "")
diagnostics_md = (
f"**Latency**: {elapsed_s:.1f}s\n\n"
f"**Candidates scored**: {result['scores_summary']['n_candidates']}\n\n"
f"**Score distribution**: "
f"median {result['scores_summary']['median']:+.3f}, "
f"min {result['scores_summary']['min']:+.3f}, "
f"max {result['scores_summary']['max']:+.3f}\n\n"
f"Domain summary (embedding query)
\n\n"
f"```\n{domain_summary}\n```\n\nTeam context (v1.4 scoring prompt)
\n\n"
f"```\n{team_ctx[:2000]}\n```\n\n