CoolFace
Apppublic

smart-models/Placebo_AI

sourceHugging Faceupdated 3mo agoView on Hugging Face
0likes
test_retrieval.py30 linesDownload Raw Back to src
1import warnings2warnings.filterwarnings('ignore')3from langchain_chroma import Chroma4from langchain_ollama import OllamaEmbeddings5 6def test_retrieval(query):7    print("Loading vector database...")8    embeddings = OllamaEmbeddings(model='nomic-embed-text:latest')9    db = Chroma(persist_directory='vector_store', embedding_function=embeddings)10 11    print(f"\nQuery: '{query}'")12    # Using similarity search with score (lower score is generally better depending on the distance metric, usually L2)13    results = db.similarity_search_with_score(query, k=3)14 15    print("\n" + "="*50)16    print("TOP 3 RETRIEVED TEXTBOOK PASSAGES")17    print("="*50)18    for i, (doc, score) in enumerate(results, 1):19        print(f"\n--- Result {i} (Distance Score: {score:.4f}) ---")20        print(f"Book: {doc.metadata.get('book_name', 'Unknown')}")21        print(f"Page: {doc.metadata.get('page_number', 'Unknown')}")22        23        # Clean up newlines and handle unicode safely for Windows terminal24        content = doc.page_content.replace('\n', ' ').encode('ascii', 'ignore').decode('ascii')25        print(f"\nContent Snippet:\n{content[:300]}...")26        print("-" * 50)27 28if __name__ == "__main__":29    test_retrieval("Differentiate Bartter syndrome from Gitelman syndrome, detailing the specific defective ion transporters in the nephron, the resulting electrolyte imbalances, and how their pathophysiology mimics specific classes of diuretics.")30