# Simplest usage codes
from sentence_transformers import SentenceTransformer
from huggingface_hub import snapshot_download
import pandas as pd
import faiss
import pickle

# 1. Download model
print("Downloading ChatMed model...")
model_path = snapshot_download(repo_id="fc28/ChatMed-RAG")

# 2. Load components
encoder = SentenceTransformer('sentence-transformers/all-mpnet-base-v2')
index = faiss.read_index(f"{model_path}/faiss_index/index.faiss")
metadata = pd.read_csv(f"{model_path}/metadata/medllm_metadata.csv")

# 3. Search function
def search(query, k=5):
    query_vec = encoder.encode([query])
    distances, indices = index.search(query_vec, k)

    results = []
    for idx in indices[0]:
        if 0 <= idx < len(metadata):
            results.append(metadata.iloc[idx])
    return results

# 4. Example usage
results = search("ChatGPT medical education applications")
for i, result in enumerate(results):
    print(f"\n{i + 1}. {result['title']}")
    print(f"   PMID: {result['pmid']}, Year: {result['year']}")
    print(f"   Abstract: {result['abstract'][:200]}...")