Velvet_Rope / app.py
Ccre's picture
Update app.py
deff59c verified
import gradio as gr
from sentence_transformers import SentenceTransformer
import torch
import json
import torch.nn.functional as F
# --- CONFIGURATE ---
# We point to your repo to ensure it's the exact same model version
MODEL_PATH = "Ccre/gemma-embedding-model"
DATABASE_FILE = "vector_database.json"
TOP_K = 5
# --- 1. Load Model and Database ONCE on startup ---
print("Loading model and database...")
model = SentenceTransformer(MODEL_PATH)
with open(DATABASE_FILE, 'r') as f:
database = json.load(f)
stored_texts = [entry['text'] for entry in database]
stored_embeddings = torch.tensor([entry['embedding'] for entry in database])
print(f"Database with {len(stored_texts)} entries loaded. Ready for queries.")
# --- 2. The Search Function ---
def search_log(query):
if not query:
return ""
query_embedding = model.encode(query, convert_to_tensor=True)
cos_scores = F.cosine_similarity(query_embedding, stored_embeddings, dim=1)
top_results = torch.topk(cos_scores, k=min(TOP_K, len(stored_texts)))
results_list = []
for score, idx in zip(top_results.values, top_results.indices):
results_list.append(f"(Score: {score:.4f}) {stored_texts[idx]}")
return "\n".join(results_list)
# --- 3. Create the Gradio Web Interface and API ---
iface = gr.Interface(
fn=search_log,
inputs=gr.Textbox(lines=2, label="Search Query"),
outputs=gr.Textbox(label="Top Results"),
title="RAG Librarian",
description="A semantic search engine for the Nova & Zanno collaboration log."
)
iface.launch(ssr_mode=False)