Repository navigation
Expand file tree
/
Copy pathrag_basic.py
More file actions
63 lines (54 loc) · 2.02 KB
/
Copy pathrag_basic.py
File metadata and controls
63 lines (54 loc) · 2.02 KB
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
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
import pypdf
from sentence_transformers import SentenceTransformer
from sklearn.metrics.pairwise import cosine_similarity
import ollama
import numpy as np
#Step 1: Load PDF and extract text
def load_pdf(file_path):
reader = pypdf.PdfReader(file_path)
text = ""
for page in reader.pages:
text += page.extract_text()
return text
#Step 2: Convert text to chuncks
def chunk_text(text, chunk_size=500):
words = text.split()
chunks = []
for i in range(0, len(words), chunk_size):
chunk = " ".join(words[i:i+chunk_size])
chunks.append(chunk)
return chunks
#Step 3: Embed Checks
def embed_chunks(chunks, model):
embeddings = model.encode(chunks)
return embeddings
# Step 4: Find relevant chunks
def get_relevant_chunks(question, chunks, embeddings, model,top_k=5):
question_embedding = model.encode([question])
similarities = cosine_similarity(question_embedding, embeddings)[0]
top_indices = np.argsort(similarities)[::-1][:top_k]
return [chunks[i] for i in top_indices]
# Step 5: Ask Ollama
def ask_ollama(question, context_chunks):
context = "\n\n".join(context_chunks)
prompt = f"""You are a healthcare VBC analyst. Answer based only on this context:
{context}
Question: {question}
Answer:"""
response = ollama.chat(model="llama3.2", messages=[{"role":"user","content":prompt}])
return response['message']['content']
#main
model = SentenceTransformer('all-MiniLM-L6-v2')
pdf_text = load_pdf("AsthmaDBR2025v110.pdf")
chunks = chunk_text(pdf_text)
print(f"Total chunks: {len(chunks)}")
embeddings = embed_chunks(chunks, model)
print("Embeddings ready.")
print(repr(pdf_text[:500]))
user_input = input("Ask a question about the asthma episode: ")
while user_input.lower() != "quit":
relevant = get_relevant_chunks(user_input, chunks, embeddings, model)
answer = ask_ollama(user_input, relevant)
print(f"\nQuestion: {user_input}")
print(f"\nAnswer: {answer}")
user_input = input("Ask a question about the asthma episode (or type 'quit' to exit): ")