File size: 3,775 Bytes
45126e4 d4d1a0c 6fd54c1 45126e4 d4d1a0c d277fdd d4d1a0c d277fdd d4d1a0c 45126e4 d4d1a0c d277fdd d4d1a0c d277fdd d4d1a0c 45126e4 d4d1a0c 6fd54c1 d277fdd d4d1a0c d277fdd d4d1a0c 6fd54c1 d277fdd d4d1a0c d277fdd d4d1a0c 6fd54c1 d277fdd d4d1a0c d277fdd d4d1a0c 6fd54c1 d277fdd d4d1a0c d277fdd d4d1a0c 45126e4 d4d1a0c 45126e4 d4d1a0c | 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 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 | import os
import faiss
import numpy as np
import gradio as gr
from sentence_transformers import SentenceTransformer
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
from PyPDF2 import PdfReader
# -----------------------------
# CONFIG
# -----------------------------
DATA_PATH = "Docs"
TOP_K = 3
# -----------------------------
# EMBEDDING MODEL (LIGHT)
# -----------------------------
embedding_model = SentenceTransformer("all-MiniLM-L6-v2")
# -----------------------------
# OPEN LLM (NO AUTH REQUIRED)
# -----------------------------
LLM_MODEL = "google/flan-t5-base"
tokenizer = AutoTokenizer.from_pretrained(LLM_MODEL)
llm_model = AutoModelForSeq2SeqLM.from_pretrained(LLM_MODEL)
# -----------------------------
# FILE LOADER
# -----------------------------
def read_file(path):
if path.endswith(".txt") or path.endswith(".md"):
with open(path, "r", encoding="utf-8") as f:
return f.read()
elif path.endswith(".pdf"):
reader = PdfReader(path)
text = ""
for page in reader.pages:
text += page.extract_text() or ""
return text
return ""
def load_docs(folder):
texts = []
for file in os.listdir(folder):
path = os.path.join(folder, file)
try:
txt = read_file(path)
if txt.strip():
texts.append(txt)
except:
continue
return texts
# -----------------------------
# CHUNKING
# -----------------------------
def chunk_text(text, size=300, overlap=50):
words = text.split()
chunks = []
for i in range(0, len(words), size - overlap):
chunks.append(" ".join(words[i:i + size]))
return chunks
# -----------------------------
# BUILD VECTOR DB
# -----------------------------
def build_index(docs):
chunks = []
for doc in docs:
chunks.extend(chunk_text(doc))
if not chunks:
return None, []
embeddings = embedding_model.encode(chunks)
dim = embeddings.shape[1]
index = faiss.IndexFlatL2(dim)
index.add(np.array(embeddings))
return index, chunks
# -----------------------------
# RETRIEVE
# -----------------------------
def retrieve(query, index, chunks, k=TOP_K):
q_embed = embedding_model.encode([query])
D, I = index.search(np.array(q_embed), k)
return [chunks[i] for i in I[0]]
# -----------------------------
# GENERATE ANSWER
# -----------------------------
def generate_answer(query, contexts):
context = "\n\n".join(contexts)
prompt = f"""
Answer the question based ONLY on the context.
If not found, say: Not in knowledge base.
Context:
{context}
Question:
{query}
"""
inputs = tokenizer(prompt, return_tensors="pt", truncation=True)
outputs = llm_model.generate(**inputs, max_new_tokens=200)
return tokenizer.decode(outputs[0], skip_special_tokens=True)
# -----------------------------
# INIT
# -----------------------------
docs = load_docs(DATA_PATH)
index, chunks = build_index(docs)
# -----------------------------
# RAG PIPELINE
# -----------------------------
def rag(query):
if index is None:
return "No documents found", ""
retrieved = retrieve(query, index, chunks)
answer = generate_answer(query, retrieved)
return answer, "\n\n---\n\n".join(retrieved)
# -----------------------------
# UI
# -----------------------------
with gr.Blocks() as demo:
gr.Markdown("## AI/ML Knowledge RAG (Stable Version)")
q = gr.Textbox(placeholder="Ask about AI tools, companies, ML...")
ans = gr.Textbox(label="Answer")
ctx = gr.Textbox(label="Context")
gr.Button("Ask").click(rag, inputs=q, outputs=[ans, ctx])
# -----------------------------
# RUN
# -----------------------------
if __name__ == "__main__":
demo.launch()
|