fb RAG

search.py

@@ -1,32 +1,24 @@
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import RagTokenizer, RagRetriever, RagTokenForGeneration, RagConfig
 from docx import Document
 from pdfminer.high_level import extract_text
-from transformers import GPT2Tokenizer
-from dataclasses import dataclass
 from typing import List
-from tqdm import tqdm
-import os 
 import pandas as pd
 import re
-from sklearn.feature_extraction.text import TfidfVectorizer
-import numpy as np
 
-tokenizer = AutoTokenizer.from_pretrained("microsoft/phi-2", trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained("microsoft/phi-2", trust_remote_code=True)
-
-EMBEDDING_SEG_LEN = 1500
-EMBEDDING_MODEL = "gpt-4" 
-
-EMBEDDING_CTX_LENGTH = 8191
-EMBEDDING_ENCODING = "cl100k_base"
-ENCODING = "gpt2"
+# Initialize RAG components
+rag_tokenizer = RagTokenizer.from_pretrained("facebook/rag-token-base")
+rag_retriever = RagRetriever.from_pretrained("facebook/rag-token-base")
+rag_token_for_generation = RagTokenForGeneration.from_pretrained("facebook/rag-token-base")
+rag_config = RagConfig.from_pretrained("facebook/rag-token-base")
 
+# Dataclass for paragraph
 @dataclass
 class Paragraph:
     page_num: int
     paragraph_num: int
     content: str
 
+# Read PDF using pdfminer
 def read_pdf_pdfminer(file_path) -> List[Paragraph]:
     text = extract_text(file_path).replace('\n', ' ').strip()
     paragraphs = batched(text, EMBEDDING_SEG_LEN)
@@ -38,6 +30,7 @@ def read_pdf_pdfminer(file_path) -> List[Paragraph]:
         paragraph_num += 1
     return paragraphs_objs
 
+# Read DOCX file
 def read_docx(file) -> List[Paragraph]:
     doc = Document(file)
     paragraphs = []
@@ -48,14 +41,17 @@ def read_docx(file) -> List[Paragraph]:
             paragraphs.append(para)
     return paragraphs
 
+# Count tokens
 def count_tokens(text, tokenizer):
     return len(tokenizer.encode(text))
 
+# Batched processing
 def batched(iterable, n):
     l = len(iterable)
     for ndx in range(0, l, n):
         yield iterable[ndx : min(ndx + n, l)]
 
+# Compute document embeddings
 def compute_doc_embeddings(df, tokenizer):
     embeddings = {}
     for index, row in tqdm(df.iterrows(), total=df.shape[0]):
@@ -64,6 +60,7 @@ def compute_doc_embeddings(df, tokenizer):
         embeddings[index] = doc_embedding
     return embeddings
 
+# Enhanced context extraction
 def enhanced_context_extraction(document, keywords, vectorizer, tfidf_scores, top_n=5):
     paragraphs = [para for para in document.split("\n") if para]
     scores = [sum([para.lower().count(keyword) * tfidf_scores[vectorizer.vocabulary_[keyword]] for keyword in keywords if keyword in para.lower()]) for para in paragraphs]
@@ -73,6 +70,7 @@ def enhanced_context_extraction(document, keywords, vectorizer, tfidf_scores, to
     
     return " ".join(relevant_paragraphs)
 
+# Targeted context extraction
 def targeted_context_extraction(document, keywords, vectorizer, tfidf_scores, top_n=5):
     paragraphs = [para for para in document.split("\n") if para]
     scores = [sum([para.lower().count(keyword) * tfidf_scores[vectorizer.vocabulary_[keyword]] for keyword in keywords]) for para in paragraphs]
@@ -82,7 +80,7 @@ def targeted_context_extraction(document, keywords, vectorizer, tfidf_scores, to
     
     return " ".join(relevant_paragraphs)
 
-
+# Extract page and clause references
 def extract_page_and_clause_references(paragraph: str) -> str:
     page_matches = re.findall(r'Page (\d+)', paragraph)
     clause_matches = re.findall(r'Clause (\d+\.\d+)', paragraph)
@@ -92,6 +90,7 @@ def extract_page_and_clause_references(paragraph: str) -> str:
     
     return f"({page_ref}, {clause_ref})".strip(", ")
 
+# Refine answer based on question
 def refine_answer_based_on_question(question: str, answer: str) -> str:
     if "Does the agreement contain" in question:
         if "not" in answer or "No" in answer:
@@ -103,7 +102,8 @@ def refine_answer_based_on_question(question: str, answer: str) -> str:
 
     return refined_answer
 
-def answer_query_with_context(question: str, df: pd.DataFrame, tokenizer, model, top_n_paragraphs: int = 5) -> str:
+# Answer query with context using RAG
+def answer_query_with_context(question: str, df: pd.DataFrame, tokenizer, retriever, generator, top_n_paragraphs: int = 5) -> str:
     question_words = set(question.split())
     
     priority_keywords = ["duration", "term", "period", "month", "year", "day", "week", "agreement", "obligation", "effective date"]
@@ -113,10 +113,13 @@ def answer_query_with_context(question: str, df: pd.DataFrame, tokenizer, model,
     most_relevant_paragraphs = df.sort_values(by='relevance_score', ascending=False).iloc[:top_n_paragraphs]['content'].tolist()
     
     context = "\n\n".join(most_relevant_paragraphs)
-    prompt = f"Question: {question}\n\nContext: {context}\n\nAnswer:"
     
-    inputs = tokenizer.encode(prompt, return_tensors="pt", max_length=512, truncation=True)
-    outputs = model.generate(inputs, max_length=200)
+    # Retrieve documents relevant to the question
+    documents = retriever.retrieve(question)
+    
+    # Generate answer using RAG
+    inputs = tokenizer(question, context, return_tensors="pt", max_length=512, truncation=True)
+    outputs = generator.generate(input_ids=inputs["input_ids"], attention_mask=inputs["attention_mask"], max_length=200, num_return_sequences=1)
     answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
     
     references = extract_page_and_clause_references(context)
@@ -124,6 +127,7 @@ def answer_query_with_context(question: str, df: pd.DataFrame, tokenizer, model,
     
     return answer
 
+# Get embedding
 def get_embedding(text, tokenizer):
     try:
         inputs = tokenizer(text, return_tensors="pt", max_length=512, truncation=True)
@@ -133,3 +137,9 @@ def get_embedding(text, tokenizer):
         print("Error obtaining embedding:", e)
         embedding = []
     return embedding
+
+# Example usage
+question = "What is the duration of the agreement?"
+df = pd.DataFrame(...)  # Assuming you have a DataFrame with content
+answer = answer_query_with_context(question, df, rag_tokenizer, rag_retriever, rag_token_for_generation)
+print("Answer:", answer)