ChatCM-RAG / ChatCM-RAG usage codes.py
fc28's picture
Upload ChatCM-RAG usage codes.py
3990efd verified
# 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]}...")