willco-afk commited on
Commit
52bbfc2
·
verified ·
1 Parent(s): 9f3d22b

Create rag_pipeline.py

Browse files
Files changed (1) hide show
  1. rag_pipeline.py +25 -0
rag_pipeline.py ADDED
@@ -0,0 +1,25 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import faiss
2
+ import pickle
3
+ from transformers import pipeline
4
+
5
+ # Load FAISS index
6
+ with open('faiss_index.index', 'rb') as f:
7
+ faiss_index = pickle.load(f)
8
+
9
+ # Load a pre-trained generative model (e.g., GPT-3 or T5)
10
+ generator = pipeline("text-generation", model="gpt2")
11
+
12
+ # Example query
13
+ query = "What is the capital of France?"
14
+
15
+ # Search for the most similar document using FAISS
16
+ query_embedding = model.encode([query])
17
+ D, I = faiss_index.search(query_embedding, k=1) # k=1 for the most similar document
18
+
19
+ # Use the retrieved document as context for the generative model
20
+ retrieved_doc = documents[I[0][0]]
21
+
22
+ # Generate a response using the retrieved document as context
23
+ prompt = f"Context: {retrieved_doc}\nQuestion: {query}\nAnswer:"
24
+ answer = generator(prompt, max_length=50)
25
+ print(answer[0]['generated_text'])