File size: 1,452 Bytes
5f0f6c4 8d78ad1 dceefee 0605cc4 8d78ad1 5f0f6c4 0605cc4 5f0f6c4 0605cc4 5f0f6c4 0605cc4 | 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 67 68 | import os
import pickle
import faiss
from sentence_transformers import SentenceTransformer
from huggingface_hub import snapshot_download
# Global variables
repo_path = None
model = None
index = None
master = None
def load_resources():
global repo_path, model, index, master
if model is None:
print("Loading resources...")
repo_path = snapshot_download(
repo_id="waseem11/master",
repo_type="model"
)
# Load SentenceTransformer
model = SentenceTransformer(
os.path.join(repo_path, "all-MiniLM-L6-v2"),
device="cpu"
)
# Load FAISS index
index = faiss.read_index(
os.path.join(repo_path, "faiss.index")
)
# Load master metadata
with open(
os.path.join(repo_path, "master.pkl"),
"rb"
) as f:
master = pickle.load(f)
print("Resources loaded successfully")
def find_company_items(query, top_k=5):
load_resources()
# Generate query embedding
embedding = model.encode(
[query],
normalize_embeddings=True
)
# Search FAISS
scores, ids = index.search(embedding, top_k)
results = []
for score, idx in zip(scores[0], ids[0]):
if idx == -1:
continue
item = master[idx].copy()
item["similarity"] = float(score)
results.append(item)
return results |