|
|
import faiss |
|
|
import numpy as np |
|
|
from fastapi import FastAPI, Query, HTTPException |
|
|
from fastapi.responses import JSONResponse |
|
|
from datasets import load_dataset |
|
|
from sentence_transformers import SentenceTransformer |
|
|
from typing import List, Dict, Tuple |
|
|
|
|
|
app = FastAPI() |
|
|
|
|
|
FIELDS = ( |
|
|
"full_name", |
|
|
"description", |
|
|
"default_branch", |
|
|
"open_issues", |
|
|
"stargazers_count", |
|
|
"forks_count", |
|
|
"watchers_count", |
|
|
"license", |
|
|
"size", |
|
|
"fork", |
|
|
"updated_at", |
|
|
"has_build_zig", |
|
|
"has_build_zig_zon", |
|
|
"created_at", |
|
|
) |
|
|
|
|
|
print("Loading sentence transformer model (all-MiniLM-L6-v2)...") |
|
|
model = SentenceTransformer("all-MiniLM-L6-v2") |
|
|
print("Model loaded successfully.") |
|
|
|
|
|
def load_and_index_dataset(name: str, include_readme: bool = False) -> Tuple[faiss.IndexFlatL2, List[Dict]]: |
|
|
try: |
|
|
print(f"Loading dataset '{name}'...") |
|
|
dataset = load_dataset(name)["train"] |
|
|
|
|
|
repo_texts = [ |
|
|
" ".join(str(x.get(field, "")) for field in FIELDS) + |
|
|
(" " + x.get("readme_content", "") if include_readme else "") + |
|
|
" " + " ".join(x.get("topics", [])) |
|
|
for x in dataset |
|
|
] |
|
|
|
|
|
if not include_readme: |
|
|
dataset = [{k: v for k, v in item.items() if k != "readme_content"} for item in dataset] |
|
|
|
|
|
print(f"Creating embeddings for {len(repo_texts)} documents in '{name}'...") |
|
|
repo_embeddings = model.encode(repo_texts, show_progress_bar=True) |
|
|
|
|
|
embedding_dim = repo_embeddings.shape[1] |
|
|
index = faiss.IndexFlatL2(embedding_dim) |
|
|
index.add(np.array(repo_embeddings, dtype=np.float32)) |
|
|
|
|
|
print(f"'{name}' dataset indexed with {index.ntotal} vectors.") |
|
|
return index, list(dataset) |
|
|
except Exception as e: |
|
|
print(f"Error loading dataset '{name}': {e}") |
|
|
raise RuntimeError(f"Dataset loading/indexing failed: {name}") |
|
|
|
|
|
indices: Dict[str, Tuple[faiss.IndexFlatL2, List[Dict]]] = {} |
|
|
|
|
|
for key, readme_flag in {"packages": True, "programs": True}.items(): |
|
|
try: |
|
|
index, data = load_and_index_dataset(f"zigistry/{key}", include_readme=readme_flag) |
|
|
indices[key] = (index, data) |
|
|
except Exception as e: |
|
|
print(f"Failed to prepare index for {key}: {e}") |
|
|
indices[key] = (None, []) |
|
|
|
|
|
def perform_search(query: str, dataset_key: str, k: int) -> List[Dict]: |
|
|
index, dataset = indices.get(dataset_key, (None, [])) |
|
|
if not index: |
|
|
raise HTTPException(status_code=500, detail=f"Index not available for {dataset_key}") |
|
|
|
|
|
try: |
|
|
query_embedding = model.encode([query]) |
|
|
distances, idxs = index.search(np.array(query_embedding, dtype=np.float32), k) |
|
|
|
|
|
results = [] |
|
|
for dist, idx in zip(distances[0], idxs[0]): |
|
|
if idx == -1: |
|
|
continue |
|
|
item = dataset[int(idx)].copy() |
|
|
item["relevance_score"] = float(1.0 - dist / 2.0) |
|
|
results.append(item) |
|
|
|
|
|
return results |
|
|
except Exception as e: |
|
|
print(f"Error during search: {e}") |
|
|
raise HTTPException(status_code=500, detail="Search failed") |
|
|
|
|
|
@app.get("/searchPackages/") |
|
|
def search_packages(q: str = Query(...), k: int = Query(10)) -> JSONResponse: |
|
|
if not q: |
|
|
raise HTTPException(status_code=400, detail="Query parameter 'q' is required.") |
|
|
results = perform_search(q, "packages", k) |
|
|
return JSONResponse(content=results, headers={"Access-Control-Allow-Origin": "*"}) |
|
|
|
|
|
@app.get("/searchPrograms/") |
|
|
def search_programs(q: str = Query(...), k: int = Query(10)) -> JSONResponse: |
|
|
if not q: |
|
|
raise HTTPException(status_code=400, detail="Query parameter 'q' is required.") |
|
|
results = perform_search(q, "programs", k) |
|
|
return JSONResponse(content=results, headers={"Access-Control-Allow-Origin": "*"}) |
|
|
|
