|
|
import faiss |
|
|
import numpy as np |
|
|
from fastapi import FastAPI |
|
|
from fastapi.responses import JSONResponse |
|
|
from datasets import load_dataset |
|
|
from sentence_transformers import SentenceTransformer |
|
|
from typing import List |
|
|
|
|
|
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): |
|
|
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) |
|
|
|
|
|
print(f"Building FAISS index for '{name}'...") |
|
|
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) |
|
|
|
|
|
indices = {} |
|
|
|
|
|
for key, readme_flag in {"packages": True, "programs": True}.items(): |
|
|
index, data = load_and_index_dataset(f"zigistry/{key}", include_readme=readme_flag) |
|
|
indices[key] = (index, data) |
|
|
|
|
|
def perform_search(query: str, dataset_key: str, k: int): |
|
|
index, dataset = indices[dataset_key] |
|
|
|
|
|
query_embedding = model.encode([query]) |
|
|
query_embedding = np.array(query_embedding, dtype=np.float32) |
|
|
|
|
|
distances, idxs = index.search(query_embedding, k) |
|
|
|
|
|
results = [] |
|
|
for dist, idx in zip(distances[0], idxs[0]): |
|
|
if idx == -1: |
|
|
continue |
|
|
|
|
|
item = dataset[int(idx)].copy() |
|
|
item['relevance_score'] = 1.0 - (dist / 2.0) |
|
|
results.append(item) |
|
|
|
|
|
return results |
|
|
|
|
|
@app.get("/searchPackages/") |
|
|
def search_packages(q: str, k: int = 10): |
|
|
results = perform_search(query=q, dataset_key="packages", k=k) |
|
|
headers = {"Access-Control-Allow-Origin": "*", "Content-Type": "application/json"} |
|
|
return JSONResponse(content=results, headers=headers) |
|
|
|
|
|
@app.get("/searchPrograms/") |
|
|
def search_programs(q: str, k: int = 10): |
|
|
results = perform_search(query=q, dataset_key="programs", k=k) |
|
|
headers = {"Access-Control-Allow-Origin": "*", "Content-Type": "application/json"} |
|
|
return JSONResponse(content=results, headers=headers) |
