|
|
import faiss |
|
|
import numpy as np |
|
|
from fastapi import FastAPI, Query |
|
|
from fastapi.responses import JSONResponse |
|
|
from datasets import load_dataset |
|
|
from sentence_transformers import SentenceTransformer |
|
|
from typing import List, Dict |
|
|
|
|
|
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", |
|
|
) |
|
|
|
|
|
model = SentenceTransformer("all-MiniLM-L6-v2") |
|
|
|
|
|
def load_dataset_with_fields(name, include_readme=False): |
|
|
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] |
|
|
return dataset, repo_texts |
|
|
|
|
|
datasets = { |
|
|
"packages": load_dataset_with_fields("zigistry/packages", include_readme=True), |
|
|
"programs": load_dataset_with_fields("zigistry/programs", include_readme=True), |
|
|
} |
|
|
|
|
|
indices = {} |
|
|
for key, (dataset, repo_texts) in datasets.items(): |
|
|
repo_embeddings = model.encode(repo_texts) |
|
|
index = faiss.IndexFlatL2(repo_embeddings.shape[1]) |
|
|
index.add(np.array(repo_embeddings)) |
|
|
indices[key] = (index, dataset) |
|
|
|
|
|
def filter_results_by_distance(distances, idxs, dataset, threshold_ratio=0.3): |
|
|
if len(distances) == 0: |
|
|
return [] |
|
|
min_distance = np.min(distances) |
|
|
max_distance = np.max(distances) |
|
|
threshold = min_distance + ((max_distance - min_distance) * threshold_ratio) |
|
|
|
|
|
results = [ |
|
|
dataset[int(i)] |
|
|
for d, i in zip(distances, idxs) |
|
|
if d <= threshold |
|
|
] |
|
|
return results |
|
|
|
|
|
@app.get("/searchPackages/") |
|
|
def search_packages(q: str): |
|
|
key = "packages" |
|
|
index, dataset = indices[key] |
|
|
query_embedding = model.encode([q]) |
|
|
distances, idxs = index.search(np.array(query_embedding), len(dataset)) |
|
|
results = filter_results_by_distance(distances[0], idxs[0], dataset) |
|
|
headers = {"Access-Control-Allow-Origin": "*", "Content-Type": "application/json"} |
|
|
return JSONResponse(content=results, headers=headers) |
|
|
|
|
|
@app.get("/searchPrograms/") |
|
|
def search_programs(q: str): |
|
|
key = "programs" |
|
|
index, dataset = indices[key] |
|
|
query_embedding = model.encode([q]) |
|
|
distances, idxs = index.search(np.array(query_embedding), len(dataset)) |
|
|
results = filter_results_by_distance(distances[0], idxs[0], dataset) |
|
|
headers = {"Access-Control-Allow-Origin": "*", "Content-Type": "application/json"} |
|
|
return JSONResponse(content=results, headers=headers) |
|
|
|
