Create rag_pipeline.py

rag_pipeline.py

@@ -0,0 +1,31 @@
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from sentence_transformers import SentenceTransformer
+import numpy as np
+import faiss
+from datasets import load_dataset
+
+# Load Dataset
+dataset = load_dataset("pubmed_qa", "pqa_labeled")
+corpus = [entry['context'] for entry in dataset['train']]
+
+# Embedding model
+embed_model = SentenceTransformer('pritamdeka/BioBERT-mnli-snli-scinli-scitail-mednli-stsb')
+corpus_embeddings = embed_model.encode(corpus, show_progress_bar=True)
+
+# FAISS index
+index = faiss.IndexFlatL2(len(corpus_embeddings[0]))
+index.add(np.array(corpus_embeddings))
+
+# Generator model
+tokenizer = AutoTokenizer.from_pretrained("facebook/bart-large")
+model = AutoModelForSeq2SeqLM.from_pretrained("facebook/bart-large")
+
+# Generate Answer Function
+def generate_answer(query, index, embeddings, corpus, embed_model):
+    query_embedding = embed_model.encode([query])
+    D, I = index.search(np.array(query_embedding), k=5)
+    retrieved = [corpus[i] for i in I[0]]
+    prompt = f"Context: {retrieved}\n\nQuestion: {query}\n\nAnswer:"
+    inputs = tokenizer(prompt, return_tensors="pt", truncation=True)
+    outputs = model.generate(**inputs, max_new_tokens=128)
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)