Views
No views yet
Document (with PMID and Title retained as metadata) and split
into 256-token chunks with a chunk overlap of chunk_size / 10 (~26 tokens) using a recursive character splitter. The dataset was
shuffled and split 80/20 into training and test sets with random_state=42, yielding 529 training and 133 test examples. To evaluate
the RAG pipeline specifically, I constructed 5 manually written test questions covering common running injuries (stress fracture prevention,
safe return-to-running, Achilles tendinopathy, IT band syndrome, and stress fracture warning signs), used to qualitatively compare retrieval
quality across embedding/distance-metric combinations before selecting a final retriever.do_sample=False, max_new_tokens=250). I compared three embedding-model/distance-metric combinations
on 5 manually-constructed test questions: all-MiniLM-L6-v2 + cosine, BAAI/bge-small-en-v1.5 + cosine, and BAAI/bge-small-en-v1.5 + Euclidean distance.
MiniLM occasionally retrieved an off-topic result (e.g., an ultra-marathon performance paper for a stress-fracture-prevention question), while bge-small
returned exclusively on-topic abstracts under either distance metric — the two bge-small variants returned identical top-3 documents for every test question,
indicating the embedding model mattered more than the distance metric for this corpus. I selected BAAI/bge-small-en-v1.5 with cosine similarity,
retrieving the top 3 chunks per query (k=3) as context for generation.wandb/RAGTruth-processed),
and PubMedQA (pqa_labeled, 1,000 expert-annotated examples) which were chosen because
all three target retrieval grounded biomedical/medicine/QA generation rather than open-domain QA, making them a closer match to the
Avid Runner task than general RAG benchmarks.long_answer) rather than the dataset's
underlying yes/no/maybe label). For comparison models, I used the base model (Qwen2.5-7B-Instruct with no retrieval) to isolate the
effect of RAG, plus Qwen2.5-1.5B-Instruct and Llama-3.2-3B-Instruct as similarly-scoped instruction-tuned models with plausible baseline
performance on the same task, drawing on the zero/3/8-shot comparison I ran earlier in the project.| Model / Condition | PubMed Test Split (ROUGE-L) | RAGBench (pubmedqa) | RAGTruth (QA) | PubMedQA (pqa_labeled) |
|---|---|---|---|---|
| Qwen2.5-7B-Instruct + RAG (this model) | 0.173 | 0.339 | 0.352 | 0.224 |
| Qwen2.5-7B-Instruct (no retrieval, base) | 0.146 | 0.198 | 0.218 | 0.124 |
| Qwen2.5-1.5B-Instruct + RAG | 0.170 | 0.270 | 0.298 | 0.207 |
| Llama-3.2-3B-Instruct + RAG | 0.174 | 0.273 | 0.365 | 0.237 |
1import os
2import torch
3from transformers import AutoTokenizer, AutoModelForCausalLM
4from langchain_community.vectorstores import FAISS
5from langchain_huggingface import HuggingFaceEmbeddings
6
7HF_TOKEN = os.environ.get("HF_TOKEN")
8
9class AvidRunnerRAGPipeline:
10 def __init__(self, model_name: str, embedding_model_name: str, vector_db_path: str):
11 self.max_new_tokens = 250
12
13 print(f"Loading Model: {model_name}...")
14 self.tokenizer = AutoTokenizer.from_pretrained(model_name, token=HF_TOKEN)
15
16 # Load to CPU first — ZeroGPU has no GPU available during startup;
17 # weights move to GPU automatically at call time.
18 self.model = AutoModelForCausalLM.from_pretrained(
19 model_name, device_map="cpu", torch_dtype=torch.bfloat16, token=HF_TOKEN
20 )
21 self.tokenizer.pad_token_id = self.tokenizer.eos_token_id
22 self.tokenizer.padding_side = "left"
23
24 print("Loading Embeddings...")
25 self.embedding_model = HuggingFaceEmbeddings(
26 model_name=embedding_model_name,
27 model_kwargs={"device": "cpu"},
28 encode_kwargs={"normalize_embeddings": True}, # matches the winning cosine-similarity retriever
29 )
30
31 print(f"Loading Vector DB from {vector_db_path}...")
32 self.vector_db = FAISS.load_local(vector_db_path, self.embedding_model, allow_dangerous_deserialization=True)
33 print("RAG Pipeline Initialized (CPU Mode)")
34
35 def retrieve(self, query, num_docs=3):
36 return self.vector_db.similarity_search(query, k=num_docs)
37
38 def _format_prompt(self, query, retrieved_docs):
39 context = "\n\n".join(
40 f"[PMID {d.metadata.get('PMID', 'N/A')}] {d.metadata.get('Title', 'Untitled')}:\n{d.page_content}"
41 for d in retrieved_docs
42 )
43 messages = [
44 {
45 "role": "system",
46 "content": (
47 "You are a sports medicine assistant helping recreational runners. "
48 "Using only the context provided below, answer the question accurately and concisely.\n\n"
49 f"Context:\n{context}"
50 ),
51 },
52 {"role": "user", "content": query},
53 ]
54 return self.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
55
56 def generate(self, query, num_docs=3):
57 retrieved_docs = self.retrieve(query, num_docs)
58 prompt_str = self._format_prompt(query, retrieved_docs)
59 inputs = self.tokenizer(prompt_str, return_tensors="pt").to(self.model.device)
60 with torch.no_grad():
61 outputs = self.model.generate(
62 **inputs, max_new_tokens=self.max_new_tokens,
63 do_sample=False, pad_token_id=self.tokenizer.eos_token_id,
64 )
65 input_len = inputs.input_ids.shape[1]
66 return self.tokenizer.decode(outputs[0][input_len:], skip_special_tokens=True).strip()
67
68
69# --- Example ---
70rag = AvidRunnerRAGPipeline(
71 model_name="Qwen/Qwen2.5-7B-Instruct",
72 embedding_model_name="BAAI/bge-small-en-v1.5",
73 vector_db_path="index/",
74)
75print(rag.generate("What advice can you give a runner regarding stress fracture prevention during marathon training?"))apply_chat_template so the same code works across backbone models. The system message sets the sports-medicine assistant role, instructs the model to answer only from the provided context, and embeds the retrieved abstract chunks directly; the user message carries the question.System:
You are a sports medicine assistant helping recreational runners.
Using only the context provided below, answer the question accurately and concisely.
Context:
{retrieved PubMed abstract chunks — top-3 by cosine similarity}
User:
{question}Question: What advice can you give a runner regarding A High Rate of Return to Running Is
Seen After Both Arthroscopic and Open Shoulder Surgery.?Based on the study, you can advise that after both arthroscopic and open shoulder surgery,
most patients are able to run 1 mile by 3 months postoperatively and about three-quarters
can do so by 4.5 months. It's important to follow the rehabilitation protocol, which
discourages running before 3 months, but once cleared by your healthcare provider, you can
aim to gradually return to running as part of your recovery process.results/pubmed_rag_results.csv):1{
2 "question": "What advice can you give a runner regarding A High Rate of Return to Running Is Seen After Both Arthroscopic and Open Shoulder Surgery.?",
3 "retrieved_sources": [
4 "A High Rate of Return to Running Is Seen After Both Arthroscopic and Open Shoulder Surgery.",
5 "Pre-operative comparisons for a return to running protocol in recreational athletes following hip arthroscopy.",
6 "A High Rate of Return to Running Is Seen After Both Arthroscopic and Open Shoulder Surgery."
7 ],
8 "response": "Based on the study, you can advise that after both arthroscopic and open shoulder surgery, most patients are able to run 1 mile by 3 months postoperatively and about three-quarters can do so by 4.5 months. It's important to follow the rehabilitation protocol, which discourages running before 3 months, but once cleared by your healthcare provider, you can aim to gradually return to running as part of your recovery process."
9}retrieved_sources — the same paper was retrieved via two different chunks. This is a real logged example, not a cleaned-up illustration, and it's the retrieval-diversity issue flagged in Limitations below.