Pythonrag_pipeline.py
Copy
1
import chromadb
2
from openai import OpenAI
3
 
4
client = OpenAI()
5
chroma = chromadb.Client()
6
collection = chroma.create_collection("docs")
7
 
8
DOCUMENTS = [
9
    ("doc1", "Transformers use self-attention to process sequences in parallel. The attention mechanism assigns weights to each token based on relevance."),
10
    ("doc2", "RAG combines retrieval with generation. It fetches relevant context from a vector store and injects it into the prompt."),
11
    ("doc3", "Fine-tuning adapts a pre-trained model to a specific task using a labeled dataset. It updates model weights unlike RAG."),
12
]
13
 
14
def embed(text: str) -> list[float]:
15
    return client.embeddings.create(model="text-embedding-3-small", input=text).data[0].embedding
16
 
17
# Index documents
18
for doc_id, text in DOCUMENTS:
19
    collection.add(ids=[doc_id], documents=[text], embeddings=[embed(text)])
20
 
21
def retrieve(query: str, top_k: int = 2) -> list[str]:
22
    results = collection.query(query_embeddings=[embed(query)], n_results=top_k)
23
    return results["documents"][0]
24
 
25
def rag_answer(question: str) -> str:
26
    chunks = retrieve(question)
27
    context = "\n---\n".join(chunks)
28
 
29
    resp = client.chat.completions.create(
30
        model="gpt-4o-mini",
31
        messages=[
32
            {"role": "system", "content": f"Answer using only the provided context.\n\nContext:\n{context}"},
33
            {"role": "user", "content": question},
34
        ],
35
    )
36
    return resp.choices[0].message.content or ""
37
 
38
print(rag_answer("How does RAG differ from fine-tuning?"))
39