File size: 2,451 Bytes
0861ec7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 |
import faiss
import numpy as np
from fastapi import FastAPI, Query
from datasets import load_dataset
from sentence_transformers import SentenceTransformer
app = FastAPI()
FIELDS = (
"full_name", "description", "watchers_count", "forks_count", "license",
"default_branch", "has_build_zig", "has_build_zig_zon", "fork",
"open_issues", "stargazers_count", "updated_at", "created_at",
"size"
)
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", "")) * include_readme +
" " + " ".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)
scroll_data = {
"infiniteScrollPackages": load_dataset_with_fields("zigistry/packages", include_readme=False)[0],
"infiniteScrollPrograms": load_dataset_with_fields("zigistry/programs", include_readme=False)[0],
}
@app.get("/fetch_data/")
def fetch_data(category: str, page_number: int = Query(0, ge=0)):
if category not in scroll_data:
return {"error": "Invalid category"}
start = page_number * 10
return scroll_data[category][start : start + 10]
@app.get("/search_repositories/")
def search_repositories(category: str, query: str):
key = "packages" if category == "SearchPackages" else "programs"
if key not in indices:
return {"error": "Invalid category"}
index, dataset = indices[key]
query_embedding = model.encode([query])
distances, indices_ = index.search(np.array(query_embedding), len(dataset))
min_distance = distances[0][0]
threshold = min_distance * 1.5
results = [dataset[int(i)] for d, i in zip(distances[0], indices_[0]) if d <= threshold]
return results[:280] if len(results) > 280 else results
|