Spaces:
Sleeping
Sleeping
import streamlit as st | |
from transformers import AutoTokenizer, AutoModelForCausalLM | |
from sentence_transformers import SentenceTransformer, util | |
import PyPDF2 | |
from docx import Document | |
import numpy as np | |
# Load the tokenizer and model for sentence embeddings | |
def load_model(): | |
try: | |
tokenizer = AutoTokenizer.from_pretrained("rakeshkiriyath/gpt2Medium_text_to_sql") | |
model = AutoModelForCausalLM.from_pretrained("rakeshkiriyath/gpt2Medium_text_to_sql") | |
sentence_model = SentenceTransformer('all-MiniLM-L6-v2') # Smaller, faster sentence embeddings model | |
st.success("Model loaded successfully!") | |
return tokenizer, model, sentence_model | |
except Exception as e: | |
st.error(f"Error loading models: {e}") | |
return None, None, None | |
# Extract text from a PDF file | |
def extract_text_from_pdf(pdf_file): | |
try: | |
pdf_reader = PyPDF2.PdfReader(pdf_file) | |
text = "" | |
for page in pdf_reader.pages: | |
text += page.extract_text() | |
return text | |
except Exception as e: | |
st.error(f"Error reading PDF: {e}") | |
return "" | |
# Extract text from a Word document | |
def extract_text_from_word(docx_file): | |
try: | |
doc = Document(docx_file) | |
text = "" | |
for paragraph in doc.paragraphs: | |
text += paragraph.text + "\n" | |
return text | |
except Exception as e: | |
st.error(f"Error reading Word document: {e}") | |
return "" | |
# Optimized comparison using embeddings and matrix operations | |
def compare_sentences(doc1_sentences, doc2_sentences, sentence_model): | |
# Encode all sentences in batches to get embeddings | |
doc1_embeddings = sentence_model.encode(doc1_sentences, convert_to_tensor=True, batch_size=16) | |
doc2_embeddings = sentence_model.encode(doc2_sentences, convert_to_tensor=True, batch_size=16) | |
# Compute cosine similarity matrix between all pairs | |
similarity_matrix = util.pytorch_cos_sim(doc1_embeddings, doc2_embeddings) | |
# Extract pairs with similarity > threshold | |
threshold = 0.6 # Adjust this for stricter or looser matching | |
similar_sentences = [] | |
for i, row in enumerate(similarity_matrix): | |
for j, score in enumerate(row): | |
if score >= threshold: | |
similar_sentences.append((i, j, score.item(), doc1_sentences[i], doc2_sentences[j])) | |
return similar_sentences | |
# Streamlit UI | |
def main(): | |
st.title("Optimized Comparative Analysis of Two Documents") | |
st.sidebar.header("Upload Files") | |
# Upload files | |
uploaded_file1 = st.sidebar.file_uploader("Upload the First Document (PDF/Word)", type=["pdf", "docx"]) | |
uploaded_file2 = st.sidebar.file_uploader("Upload the Second Document (PDF/Word)", type=["pdf", "docx"]) | |
if uploaded_file1 and uploaded_file2: | |
# Extract text from the uploaded documents | |
if uploaded_file1.name.endswith(".pdf"): | |
text1 = extract_text_from_pdf(uploaded_file1) | |
else: | |
text1 = extract_text_from_word(uploaded_file1) | |
if uploaded_file2.name.endswith(".pdf"): | |
text2 = extract_text_from_pdf(uploaded_file2) | |
else: | |
text2 = extract_text_from_word(uploaded_file2) | |
if not text1.strip(): | |
st.error("The first document is empty or could not be read.") | |
return | |
if not text2.strip(): | |
st.error("The second document is empty or could not be read.") | |
return | |
st.write("### Preview of Document 1:") | |
st.text(text1[:500]) # Display a preview of Document 1 | |
st.write("### Preview of Document 2:") | |
st.text(text2[:500]) # Display a preview of Document 2 | |
# Split text into sentences | |
doc1_sentences = text1.split('. ') | |
doc2_sentences = text2.split('. ') | |
# Limit sentences for testing purposes (optional) | |
doc1_sentences = doc1_sentences[:50] # Remove this line for full processing | |
doc2_sentences = doc2_sentences[:50] # Remove this line for full processing | |
# Load models | |
tokenizer, model, sentence_model = load_model() | |
if not sentence_model: | |
st.error("Failed to load the sentence embedding model.") | |
return | |
# Perform sentence comparison | |
st.info("Comparing sentences, this may take a moment...") | |
similar_sentences = compare_sentences(doc1_sentences, doc2_sentences, sentence_model) | |
# Display results | |
st.header("Comparative Analysis Results") | |
st.write(f"Number of sentences in Document 1: {len(doc1_sentences)}") | |
st.write(f"Number of sentences in Document 2: {len(doc2_sentences)}") | |
if similar_sentences: | |
st.success(f"Found {len(similar_sentences)} similar sentences!") | |
for match in similar_sentences: | |
doc1_index, doc2_index, score, sent1, sent2 = match | |
st.markdown(f"**Document 1 Sentence {doc1_index + 1}:** {sent1}") | |
st.markdown(f"**Document 2 Sentence {doc2_index + 1}:** {sent2}") | |
st.markdown(f"**Similarity Score:** {score:.2f}") | |
st.markdown("---") | |
else: | |
st.info("No significantly similar sentences found.") | |
else: | |
st.warning("Please upload two documents to compare.") | |
if __name__ == "__main__": | |
main() | |