File size: 5,209 Bytes
a6fe316 ade8cf8 90d214f 378f418 d0eefa8 378f418 90d214f a6fe316 e22de9e 378f418 21583be 378f418 21583be 378f418 a6fe316 378f418 a6fe316 378f418 a6fe316 90d214f 378f418 90d214f 378f418 90d214f 378f418 350d57a 378f418 90d214f 378f418 90d214f 378f418 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 |
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 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"
@dataclass
class Paragraph:
page_num: int
paragraph_num: int
content: str
def read_pdf_pdfminer(file_path) -> List[Paragraph]:
text = extract_text(file_path).replace('\n', ' ').strip()
paragraphs = batched(text, EMBEDDING_SEG_LEN)
paragraphs_objs = []
paragraph_num = 1
for p in paragraphs:
para = Paragraph(0, paragraph_num, p)
paragraphs_objs.append(para)
paragraph_num += 1
return paragraphs_objs
def read_docx(file) -> List[Paragraph]:
doc = Document(file)
paragraphs = []
for paragraph_num, paragraph in enumerate(doc.paragraphs, start=1):
content = paragraph.text.strip()
if content:
para = Paragraph(1, paragraph_num, content)
paragraphs.append(para)
return paragraphs
def count_tokens(text, tokenizer):
return len(tokenizer.encode(text))
def batched(iterable, n):
l = len(iterable)
for ndx in range(0, l, n):
yield iterable[ndx : min(ndx + n, l)]
def compute_doc_embeddings(df, tokenizer):
embeddings = {}
for index, row in tqdm(df.iterrows(), total=df.shape[0]):
doc = row["content"]
doc_embedding = get_embedding(doc, tokenizer)
embeddings[index] = doc_embedding
return embeddings
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]
top_indices = sorted(range(len(scores)), key=lambda i: scores[i], reverse=True)[:top_n]
relevant_paragraphs = [paragraphs[i] for i in top_indices]
return " ".join(relevant_paragraphs)
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]
top_indices = sorted(range(len(scores)), key=lambda i: scores[i], reverse=True)[:top_n]
relevant_paragraphs = [paragraphs[i] for i in top_indices]
return " ".join(relevant_paragraphs)
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)
page_ref = f"Page {page_matches[0]}" if page_matches else ""
clause_ref = f"Clause {clause_matches[0]}" if clause_matches else ""
return f"({page_ref}, {clause_ref})".strip(", ")
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:
refined_answer = f"No, the agreement does not contain {answer}"
else:
refined_answer = f"Yes, the agreement contains {answer}"
else:
refined_answer = answer
return refined_answer
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"]
df['relevance_score'] = df['content'].apply(lambda x: len(question_words.intersection(set(x.split()))) + sum([x.lower().count(pk) for pk in priority_keywords]))
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=600)
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
def get_embedding(text, tokenizer):
try:
inputs = tokenizer(text, return_tensors="pt", max_length=512, truncation=True)
outputs = model(**inputs)
embedding = outputs.last_hidden_state
except Exception as e:
print("Error obtaining embedding:", e)
embedding = []
return embedding
|