hwca96/CS_Paper_Abstract_Semantic_Search

1from pymongo import MongoClient2from dotenv import load_dotenv3import requests4import os5import pandas as pd6 7load_dotenv()8 9class MongoSearch:10    def __init__(self, database = "arxiv_data", collection = "cs_sample"):11        self.client = MongoClient(os.getenv("MONGODB_URI"))12        self.db = self.client[database]13        self.collection = self.db[collection]14        self.embedding_url = "https://api-inference.huggingface.co/pipeline/feature-extraction/sentence-transformers/all-MiniLM-L6-v2"15 16    def generate_embedding(self, text):17        response = requests.post(18            self.embedding_url,19            headers={"Authorization": f"Bearer {os.getenv('HF_TOKEN')}"},20            json={"inputs": text, "options": {"wait_for_model": True}})21 22        if response.status_code != 200:23            raise ValueError(f"Request failed with status code {response.status_code}: {response.text}")24 25        return response.json()26 27    def search(self, query):28        results = self.collection.aggregate([29        {"$vectorSearch": {30            "queryVector": self.generate_embedding(query),31            "path": "embedding",32            "numCandidates": 100,33            "limit": 5,34            "index": "EmbeddingSemanticSearch",35            }}36        ])37        keys = ["title", "abstract", "authors", "id", "categories", "update_date"]38        return pd.DataFrame([dict((k, result[k]) for k in keys) for result in results])39