Update app.py

app.py

@@ -8,74 +8,79 @@ from typing import List, Dict
 
 app = FastAPI()
 
-# Constants
-EMBEDDING_MODEL_NAME = "all-MiniLM-L6-v2"
-SEARCH_THRESHOLD_RATIO = 0.3
-FIELDS_TO_INCLUDE = [
-    "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"
-]
+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",
+)
 
-# Load embedding model
-model = SentenceTransformer(EMBEDDING_MODEL_NAME)
+model = SentenceTransformer("all-MiniLM-L6-v2")
 
-# Helper functions
-def prepare_text(entry: Dict, include_readme: bool = True) -> str:
-    parts = [str(entry.get(field, "")) for field in FIELDS_TO_INCLUDE]
-    if include_readme:
-        parts.append(entry.get("readme_content", ""))
-    parts.extend(entry.get("topics", []))
-    return " ".join(parts)
+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
 
-def load_and_encode_dataset(name: str, include_readme: bool = True):
-    raw_dataset = load_dataset(name)["train"]
-    texts = [prepare_text(item, include_readme) for item in raw_dataset]
-    embeddings = model.encode(texts)
-    return raw_dataset, np.array(embeddings)
-
-def build_faiss_index(embeddings: np.ndarray) -> faiss.IndexFlatL2:
-    index = faiss.IndexFlatL2(embeddings.shape[1])
-    index.add(embeddings)
-    return index
+datasets = {
+    "packages": load_dataset_with_fields("zigistry/packages", include_readme=True),
+    "programs": load_dataset_with_fields("zigistry/programs", include_readme=True),
+}
 
-def search_index(index: faiss.IndexFlatL2, query: str, embeddings: np.ndarray, dataset: List[Dict]) -> List[Dict]:
-    query_vector = model.encode([query])
-    distances, indices = index.search(np.array(query_vector), len(dataset))
-    return filter_by_distance(distances[0], indices[0], dataset)
+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_by_distance(distances: np.ndarray, indices: np.ndarray, dataset: List[Dict], ratio: float = SEARCH_THRESHOLD_RATIO) -> List[Dict]:
+def filter_results_by_distance(distances, idxs, dataset, threshold_ratio=0.3):
     if len(distances) == 0:
         return []
-    min_d, max_d = np.min(distances), np.max(distances)
-    threshold = min_d + (max_d - min_d) * ratio
-    return [dataset[i] for d, i in zip(distances, indices) if d <= threshold]
-
-# Load datasets and create indices
-data_configs = {
-    "packages": "zigistry/packages",
-    "programs": "zigistry/programs"
-}
-
-data_store = {}
+    min_distance = np.min(distances)
+    max_distance = np.max(distances)
+    threshold = min_distance + ((max_distance - min_distance) * threshold_ratio)
 
-for key, dataset_name in data_configs.items():
-    dataset, embeddings = load_and_encode_dataset(dataset_name, include_readme=True)
-    index = build_faiss_index(embeddings)
-    data_store[key] = {
-        "dataset": dataset,
-        "index": index,
-        "embeddings": embeddings
-    }
+    results = [
+        dataset[int(i)]
+        for d, i in zip(distances, idxs)
+        if d <= threshold
+    ]
+    return results
 
-# FastAPI endpoints
-@app.get("/search/{category}/")
-def search(category: str, q: str = Query(...)):
-    if category not in data_store:
-        return JSONResponse(status_code=404, content={"error": "Invalid category"})
+@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)
 
-    store = data_store[category]
-    results = search_index(store["index"], q, store["embeddings"], store["dataset"])
+@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)
+    return JSONResponse(content=results, headers=headers)