Update search.py

search.py

@@ -1,32 +1,32 @@
-from transformers import RagTokenizer, RagRetriever, RagTokenForGeneration, RagConfig
+from transformers import AutoTokenizer, AutoModelForCausalLM
 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 datasets import load_dataset
+from sklearn.feature_extraction.text import TfidfVectorizer
+import numpy as np
 
-# Initialize RAG components
-rag_tokenizer = RagTokenizer.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")
+tokenizer = AutoTokenizer.from_pretrained("impira/layoutlm-document-qa", trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained("impira/layoutlm-document-qa", trust_remote_code=True)
 
-# Download and prepare the wiki_dpr dataset
-dpr_dataset = load_dataset("wiki_dpr")
-passages = dpr_dataset["train"]["passage"]
-titles = dpr_dataset["train"]["title"]
+EMBEDDING_SEG_LEN = 1500
+EMBEDDING_MODEL = "gpt-4" 
 
-# Initialize the RagRetriever
-rag_retriever = RagRetriever(passages=passages, titles=titles, config=rag_config)
+EMBEDDING_CTX_LENGTH = 8191
+EMBEDDING_ENCODING = "cl100k_base"
+ENCODING = "gpt2"
 
-# 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,7 +38,6 @@ 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 = []
@@ -49,17 +48,14 @@ 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]):
@@ -68,7 +64,6 @@ 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]
@@ -78,7 +73,6 @@ 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]
@@ -88,7 +82,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)
@@ -98,7 +92,6 @@ 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:
@@ -110,8 +103,7 @@ def refine_answer_based_on_question(question: str, answer: str) -> str:
 
     return refined_answer
 
-# 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:
+def answer_query_with_context(question: str, df: pd.DataFrame, tokenizer, model, top_n_paragraphs: int = 5) -> str:
     question_words = set(question.split())
     
     priority_keywords = ["duration", "term", "period", "month", "year", "day", "week", "agreement", "obligation", "effective date"]
@@ -121,18 +113,17 @@ def answer_query_with_context(question: str, df: pd.DataFrame, tokenizer, retrie
     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:"
     
-    # Generate answer using RAG
-    inputs = rag_tokenizer(question, context, return_tensors="pt", max_length=512, truncation=True)
-    outputs = rag_token_for_generation.generate(input_ids=inputs["input_ids"], attention_mask=inputs["attention_mask"], max_length=200, num_return_sequences=1)
-    answer = rag_tokenizer.decode(outputs[0], skip_special_tokens=True)
+    inputs = tokenizer.encode(prompt, return_tensors="pt", max_length=512, truncation=True)
+    outputs = model.generate(inputs, max_length=200)
+    answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
     
     references = extract_page_and_clause_references(context)
     answer = refine_answer_based_on_question(question, answer) + " " + references
     
     return answer
 
-# Get embedding
 def get_embedding(text, tokenizer):
     try:
         inputs = tokenizer(text, return_tensors="pt", max_length=512, truncation=True)
@@ -142,9 +133,3 @@ 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)