| import ollama |
| import numpy as np |
|
|
| |
| EMBEDDING_MODEL = "hf.co/CompendiumLabs/bge-base-en-v1.5-gguf" |
| LANGUAGE_MODEL = "hf.co/bartowski/Llama-3.2-1B-Instruct-GGUF" |
|
|
| VECTOR_DB = [] |
|
|
|
|
| def load_dataset(file_path): |
| """ |
| Load text file and return chunks. |
| """ |
| with open(file_path, "r", encoding="utf-8") as file: |
| chunks = file.readlines() |
|
|
| print(f"Loaded {len(chunks)} entries") |
| return chunks |
|
|
|
|
| def get_embedding(text): |
| response = ollama.embed( |
| model=EMBEDDING_MODEL, |
| input=text |
| ) |
| return response["embeddings"][0] |
|
|
|
|
| def build_vector_db(dataset): |
| """ |
| Create embeddings for all chunks. |
| """ |
| global VECTOR_DB |
|
|
| VECTOR_DB = [] |
|
|
| for i, chunk in enumerate(dataset): |
| embedding = get_embedding(chunk) |
|
|
| VECTOR_DB.append( |
| { |
| "text": chunk.strip(), |
| "embedding": embedding |
| } |
| ) |
|
|
|
|
| return VECTOR_DB |
|
|
|
|
| def cosine_similarity(a, b): |
| a = np.array(a) |
| b = np.array(b) |
|
|
| return np.dot(a, b) / ( |
| np.linalg.norm(a) * np.linalg.norm(b) |
| ) |
|
|
|
|
| def retrieve(query, top_n=3): |
| """ |
| Retrieve most relevant chunks. |
| """ |
|
|
| query_embedding = get_embedding(query) |
|
|
| similarities = [] |
|
|
| for item in VECTOR_DB: |
| similarity = cosine_similarity( |
| query_embedding, |
| item["embedding"] |
| ) |
|
|
| similarities.append( |
| ( |
| item["text"], |
| similarity |
| ) |
| ) |
|
|
| similarities.sort( |
| key=lambda x: x[1], |
| reverse=True |
| ) |
|
|
| return similarities[:top_n] |
|
|
|
|
| def generate_response(query, top_n=3): |
| """ |
| Full RAG pipeline. |
| """ |
|
|
| retrieved_docs = retrieve(query, top_n) |
|
|
| context = "\n".join( |
| [doc for doc, _ in retrieved_docs] |
| ) |
|
|
| system_prompt = f""" |
| You are a helpful AI assistant. |
| |
| Answer ONLY from the provided context. |
| |
| If the answer is not present in the context, |
| say: |
| "I couldn't find that information in the provided documents." |
| |
| Context: |
| {context} |
| """ |
|
|
| response = ollama.chat( |
| model=LANGUAGE_MODEL, |
| messages=[ |
| { |
| "role": "system", |
| "content": system_prompt |
| }, |
| { |
| "role": "user", |
| "content": query |
| } |
| ] |
| ) |
|
|
| return ( |
| response["message"]["content"], |
| retrieved_docs |
| ) |
|
|
|
|
| if __name__ == "__main__": |
|
|
| dataset = load_dataset("cat-facts.txt") |
|
|
| build_vector_db(dataset) |
|
|
| while True: |
| query = input("\nAsk a question: ") |
|
|
| if query.lower() == "exit": |
| break |
|
|
| answer, docs = generate_response(query) |
|
|
| print("\nRetrieved Documents:\n") |
|
|
| for doc, score in docs: |
| print(f"({score:.4f}) {doc}") |
|
|
| print("\nAnswer:\n") |
| print(answer) |
