File size: 1,248 Bytes
cbc783a ae93502 cbc783a |
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 |
from qdrant_client import QdrantClient
from sentence_transformers import SentenceTransformer
import os
from dotenv import load_dotenv
load_dotenv()
# Initialize Qdrant Cloud client
url=os.environ.get("QDRANT_URL")
api_key=os.environ.get("QDRANT_API_KEY")
qdrant_client = QdrantClient(
url=url, # Use environment variable
api_key=api_key, # Use environment variable
timeout=30.0
)
# Initialize SentenceTransformer model
cache_dir = os.environ.get("MODEL_CACHE_DIR", "/app/cache") # Fallback to /app/cache
os.makedirs(cache_dir, exist_ok=True)
model = SentenceTransformer("all-MiniLM-L6-v2", cache_folder=cache_dir)
# Collection name (must match the one used during upsert)
collection_name = "symptoms"
def get_similar_symptoms(symptom):
# Encode the input symptom
vector = model.encode(symptom)
# Perform similarity search
results = qdrant_client.search(
collection_name=collection_name,
query_vector=vector.tolist(),
limit=5, # Return top 5 similar symptoms
with_payload=True # Include payload (e.g., symptom name)
)
# Extract symptom names from results
similar_symptoms = [result.payload["name"] for result in results]
return similar_symptoms
|