Spaces:
Sleeping
Sleeping
| import json | |
| import pickle | |
| import numpy as np | |
| import faiss | |
| from fastapi import FastAPI | |
| from fastapi.middleware.cors import CORSMiddleware | |
| from pydantic import BaseModel | |
| from sentence_transformers import SentenceTransformer | |
| import os | |
| import requests | |
| from bs4 import BeautifulSoup | |
| app = FastAPI(title="SHL Assessment Recommender") | |
| app.add_middleware( | |
| CORSMiddleware, | |
| allow_origins=["*"], | |
| allow_methods=["*"], | |
| allow_headers=["*"] | |
| ) | |
| print("Loading models...") | |
| embedding_model = SentenceTransformer("all-MiniLM-L6-v2") | |
| faiss_index = faiss.read_index("vector_index.faiss") | |
| with open("assessment_data.pkl", "rb") as f: | |
| stored_assessments = pickle.load(f) | |
| print("Models loaded!") | |
| class QueryInput(BaseModel): | |
| query: str | |
| def find_similar(query, top_k=10): | |
| vec = embedding_model.encode([query]).astype("float32") | |
| _, indices = faiss_index.search(vec, top_k) | |
| results = [] | |
| for idx in indices[0]: | |
| if idx < len(stored_assessments): | |
| results.append(stored_assessments[idx]) | |
| return results | |
| def balance(candidates): | |
| p_types = ["P", "B"] | |
| k_types = ["K", "A"] | |
| p_items = [i for i in candidates if any(t in p_types for t in i.get("test_type", []))] | |
| k_items = [i for i in candidates if any(t in k_types for t in i.get("test_type", []))] | |
| o_items = [i for i in candidates if i not in p_items and i not in k_items] | |
| result = [] | |
| result.extend(p_items[:3]) | |
| result.extend(k_items[:3]) | |
| result.extend(o_items[:4]) | |
| remaining = [i for i in candidates if i not in result] | |
| while len(result) < 10 and remaining: | |
| result.append(remaining.pop(0)) | |
| return result[:10] | |
| def health(): | |
| return {"status": "healthy"} | |
| def recommend(user_input: QueryInput): | |
| query = user_input.query | |
| if query.startswith("http"): | |
| try: | |
| r = requests.get(query, timeout=10) | |
| soup = BeautifulSoup(r.content, "html.parser") | |
| query = soup.get_text(separator=" ", strip=True)[:3000] | |
| except: | |
| pass | |
| candidates = find_similar(query, top_k=20) | |
| final = balance(candidates) | |
| response = [] | |
| for item in final: | |
| response.append({ | |
| "url": item.get("url", ""), | |
| "name": item.get("name", ""), | |
| "adaptive_support": item.get("adaptive_support", "No"), | |
| "description": item.get("description", ""), | |
| "duration": item.get("duration", 0), | |
| "remote_support": item.get("remote_support", "No"), | |
| "test_type": item.get("test_type", []) | |
| }) | |
| return {"recommended_assessments": response} | |
| if __name__ == "__main__": | |
| import uvicorn | |
| port = int(os.environ.get("PORT", 8000)) | |
| uvicorn.run(app, host="0.0.0.0", port=port) | |