Upload 2 files

app.py

@@ -0,0 +1,133 @@
+
+import PyPDF2
+import faiss
+import numpy as np
+
+#Specify file paths
+file_path1 = '/content/AST-1.pdf'
+file_path2 = '/content/AST-2.pdf'
+
+#Step 1 : Load the document files
+def read_pdf(file_path):
+    with open(file_path, 'rb') as file:
+        reader = PyPDF2.PdfReader(file)
+        text = ''
+        for page_num in range(len(reader.pages)):
+            page = reader.pages[page_num]
+            text += page.extract_text()
+        return text
+
+ast1_text = read_pdf(file_path1)
+ast2_text = read_pdf(file_path2)
+
+#Step 2 Split the loaded documents into chunks
+# Split by Fixed Number of Words:
+def chunk_text(text, chunk_size=200):
+    words = text.split()
+    chunks = [' '.join(words[i:i + chunk_size]) for i in range(0, len(words), chunk_size)]
+    return chunks
+
+ast1_chunks = chunk_text(ast1_text, chunk_size=100)
+ast2_chunks = chunk_text(ast2_text, chunk_size=150)
+
+#label the chunks
+ast1_chunks = [(chunk, 'AST-1') for chunk in ast1_chunks]
+ast2_chunks = [(chunk, 'AST-2') for chunk in ast2_chunks]
+all_chunks = ast1_chunks + ast2_chunks
+
+#Load the Embedding Model and LLM
+from sentence_transformers import SentenceTransformer
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+
+# Load the pre-trained model from the MTEB leaderboard
+embedding_model  = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
+
+embeddings = embedding_model.encode(all_chunks, convert_to_tensor=True)
+
+# Convert embeddings to numpy arrays
+embeddings_np = np.array([embedding.cpu().numpy() for embedding in embeddings])
+
+# Create a FAISS index
+dimension = embeddings_np.shape[1]
+index = faiss.IndexFlatL2(dimension)
+index.add(embeddings_np)
+
+# Save the index
+faiss.write_index(index, 'embeddings_index.faiss')
+
+# Load FAISS index
+stored_index = faiss.read_index('/content/embeddings_index.faiss')
+
+#Function to retrieve chunks
+def retrieve_chunks(query, top_k=10, use_mmr=False, diversity=0.5, target_doc='AST-1'):
+    query_embedding = embedding_model.encode(query, convert_to_tensor=True).cpu().numpy()
+    distances, indices = stored_index.search(np.array([query_embedding]), top_k)
+
+    if use_mmr:
+        # Implement MMR-based retrieval
+        from sklearn.metrics.pairwise import cosine_similarity
+
+        selected_indices = []
+        selected_distances = []
+        candidate_indices = [i for i in indices[0] if all_chunks[i][1] == target_doc]
+        candidate_distances = [distances[0][i] for i in range(len(indices[0])) if all_chunks[indices[0][i]][1] == target_doc]
+
+        while len(selected_indices) < top_k and candidate_indices:
+            if not selected_indices:
+                selected_indices.append(candidate_indices.pop(0))
+                selected_distances.append(candidate_distances.pop(0))
+            else:
+                remaining_candidates = [candidate_indices[i] for i in range(len(candidate_indices))]
+                remaining_embeddings = embeddings_np[remaining_candidates]
+                selected_embeddings = embeddings_np[selected_indices]
+
+                similarities = cosine_similarity(remaining_embeddings, selected_embeddings)
+                mmr_scores = (1 - diversity) * np.array(candidate_distances[:len(remaining_candidates)]) - diversity * np.max(similarities, axis=1)
+
+                next_index = np.argmax(mmr_scores)
+                selected_indices.append(candidate_indices.pop(next_index))
+                selected_distances.append(candidate_distances.pop(next_index))
+
+        return [all_chunks[i][0] for i in selected_indices]
+    else:
+      retrieved_chunks = []
+      for idx in indices[0]:
+          chunk, doc_label = all_chunks[idx]
+          if doc_label == target_doc:
+              retrieved_chunks.append(chunk)
+          if len(retrieved_chunks) >= top_k:
+              break
+      return retrieved_chunks
+
+
+# Generate response
+def generate_response(query, retrieved_chunks):
+    #context = " ".join([chunk for chunk, _ in retrieved_chunks])
+    context = " ".join(retrieved_chunks) # for retrieved_chunks as array of strings
+    input_text = f"Query: {query}\nContext: {context}"
+    inputs = tokenizer(input_text, return_tensors='pt', max_length=1024, truncation=True)
+    summary_ids = generation_model.generate(inputs['input_ids'], max_length=150, min_length=40, length_penalty=2.0, num_beams=4, early_stopping=True)
+    return tokenizer.decode(summary_ids[0], skip_special_tokens=True)
+
+def rag_system(query, use_mmr=False):
+    retrieved_chunks = retrieve_chunks(query, top_k=3, use_mmr=use_mmr)
+    response = generate_response(query, retrieved_chunks)
+    return response
+    
+import gradio as gr
+
+def query_rag_system(query, use_mmr):
+    return rag_system(query, use_mmr=use_mmr)
+
+interface = gr.Interface(
+    fn=query_rag_system,
+    inputs=[
+        gr.Textbox(lines=2, placeholder="Enter your query here..."),
+        gr.Checkbox(label="Use MMR")
+    ],
+    outputs="text",
+    title="RAG System",
+    description="Enter a query to get a response from the RAG system. Optionally, use MMR for better results."
+)
+
+interface.launch()

requirements.txt

@@ -0,0 +1,5 @@
+PyPDF2
+gradio
+transformers
+sentence-transformers
+faiss-cpu