Update app.py

app.py

@@ -6,27 +6,69 @@ import numpy as np
 from io import StringIO
 import sys
 import time
+import pandas as pd
 from pymongo import MongoClient
+import plotly.express as px
+from pinecone import Pinecone, ServerlessSpec
+import chromadb
+import requests
+from io import BytesIO
+from PyPDF2 import PdfReader
+import hashlib
+import os
 
 # File Imports
-from embedding import get_embeddings  # Ensure this file/module is available
+from embedding import get_embeddings,get_image_embeddings,get_embed_chroma,imporve_text  # Ensure this file/module is available
 from preprocess import filtering  # Ensure this file/module is available
 from search import *
 
 
-# Mongo Connections
-srv_connection_uri = "mongodb+srv://adityasm1410:uOh6i11AYFeKp4wd@patseer.5xilhld.mongodb.net/?retryWrites=true&w=majority&appName=Patseer"
+# Chroma Connections
+client = chromadb.PersistentClient(path = "embeddings")
+collection = client.get_or_create_collection(name="data",metadata={"hnsw:space": "l2"})
 
-client = MongoClient(srv_connection_uri)
-db = client['embeddings'] 
-collection = db['data']  
+
+def generate_hash(content):
+    return hashlib.sha256(content.encode('utf-8')).hexdigest()
+
+def get_key(link):
+    text = ''
+    try:
+        # Fetch the PDF file from the URL
+        response = requests.get(link)
+        response.raise_for_status()  # Raise an error for bad status codes
+
+        # Use BytesIO to handle the PDF content in memory
+        pdf_file = BytesIO(response.content)
+
+        # Load the PDF file
+        reader = PdfReader(pdf_file)
+        num_pages = len(reader.pages)
+        
+        first_page_text = reader.pages[0].extract_text()
+        if first_page_text:
+            text += first_page_text
+        
+
+        last_page_text = reader.pages[-1].extract_text()
+        if last_page_text:
+            text += last_page_text
+
+    except requests.exceptions.HTTPError as e:
+        print(f'HTTP error occurred: {e}')
+    except Exception as e:
+        print(f'An error occurred: {e}')
+    
+    unique_key = generate_hash(text)
+    
+    return unique_key
 
 # Cosine Similarity Function
 def cosine_similarity(vec1, vec2):
     vec1 = np.array(vec1)
     vec2 = np.array(vec2)
     
-    dot_product = np.dot(vec1, vec2)
+    dot_product = np.dot(vec1, vec2.T)
     magnitude_vec1 = np.linalg.norm(vec1)
     magnitude_vec2 = np.linalg.norm(vec2)
     
@@ -36,6 +78,29 @@ def cosine_similarity(vec1, vec2):
     cosine_sim = dot_product / (magnitude_vec1 * magnitude_vec2)
     return cosine_sim
 
+def update_chroma(product_name,url,key,text,vector,log_area):
+
+    id_list = [key+str(i) for i in range(len(text))]
+
+    metadata_list = [
+            {   'key':key,
+                'product_name': product_name,
+                'url': url,
+                'text':item
+            }
+            for item in text
+        ]
+
+    collection.upsert(
+        ids = id_list,
+        embeddings = vector,
+        metadatas = metadata_list
+    )
+
+    logger.write(f"\n\u2713 Updated DB - {url}\n\n")
+    log_area.text(logger.getvalue())
+
+
 # Logger class to capture output
 class StreamCapture:
     def __init__(self):
@@ -52,12 +117,11 @@ class StreamCapture:
 # Main Function
 def score(main_product, main_url, product_count, link_count, search, logger, log_area):
 
-    existing_products_urls = set(collection.distinct('url'))
 
     data = {}
     similar_products = extract_similar_products(main_product)[:product_count]
 
-    
+    print("--> Fetching Manual Links")
     # Normal Filtering + Embedding  -----------------------------------------------
     if search == 'All':
 
@@ -107,94 +171,69 @@ def score(main_product, main_url, product_count, link_count, search, logger, log
 
 
     # Filtered Link -----------------------------------------
-    logger.write("\n\nFiltered Links ------------------>\n")
-    logger.write(str(data) + "\n")
+    logger.write("\n\n\u2713 Filtered Links\n")
     log_area.text(logger.getvalue())
 
 
-
     # Main product Embeddings ---------------------------------
-    logger.write("\n\nCreating Main product Embeddings ---------->\n")
-
-    # Check main product in MongoDB
-    if main_url in existing_products_urls:
-        saved_data = collection.find_one({'url': main_url})
+    logger.write("\n\n--> Creating Main product Embeddings\n")
 
-        if tag_option not in saved_data:
-            main_result , main_embedding = get_embeddings(main_url,tag_option)
-        else:
-            main_embedding = saved_data[tag_option]
-    else:
-        main_result , main_embedding = get_embeddings(main_url,tag_option)
+    main_key = get_key(main_url)
+    main_text,main_vector = get_embed_chroma(main_url)
 
-    log_area.text(logger.getvalue())
-    print("main",main_embedding)
+    update_chroma(main_product,main_url,main_key,main_text,main_vector,log_area)
 
-    update_doc = {
-        '$set': {
-            'product_name': main_product,
-            'url': main_url,
-            tag_option: main_embedding
-            }
-    }
+    # log_area.text(logger.getvalue())
+    print("\n\n\u2713 Main Product embeddings Created")
 
-    collection.update_one(
-        {'url': main_url},
-        update_doc,
-        upsert=True
-    )
 
+    logger.write("\n\n--> Creating Similar product Embeddings\n")
+    log_area.text(logger.getvalue()) 
+    test_embedding = [0]*768 
 
-    #Similar Products Check            
-    cosine_sim_scores = []
+    for product in data:
+        for link in data[product]:
 
-    logger.write("\n\nCreating Similar product Embeddings ---------->\n")
-    log_area.text(logger.getvalue())
+            url, _ = link
+            similar_key = get_key(url)
 
+            res = collection.query(
+                    query_embeddings = [test_embedding],
+                    n_results=1,
+                    where={"key": similar_key},
+                )
 
-    for product in data:
+            if not res['distances'][0]:
+                similar_text,similar_vector = get_embed_chroma(url)
+                update_chroma(product,url,similar_key,similar_text,similar_vector,log_area)
 
-        if len(data[product])==0:
-            logger.write("\n\nNo Product links Found Increase No of Links or Change Search Source\n")
-            log_area.text(logger.getvalue())
-    
-            cosine_sim_scores.append((product,'No Product links Found Increase Number of Links or Change Search Source',None,None))
-        
-        else:
-            for link,present in data[product][:link_count]:
-                
-                saved_data = collection.find_one({'url': link})
 
-                if present and (tag_option in saved_data):
-                    similar_embedding = saved_data[tag_option]
-                else:
-                    similar_result, similar_embedding = get_embeddings(link,tag_option)
+    logger.write("\n\n\u2713 Similar Product embeddings Created\n")
+    log_area.text(logger.getvalue())
 
-                log_area.text(logger.getvalue())
+    top_similar = []
 
-                print(similar_embedding)
-                for i in range(len(main_embedding)):
-                    score = cosine_similarity(main_embedding[i], similar_embedding[i])
-                    cosine_sim_scores.append((product, link, i, score))
-                    log_area.text(logger.getvalue())
-                
-                update_doc = {
-                    '$set': {
-                        'product_name': product,
-                        'url': link,
-                        tag_option: similar_embedding
-                    }
-                }
-
-                collection.update_one(
-                    {'url': link},
-                    update_doc,
-                    upsert=True
+    for idx,chunk in enumerate(main_vector):
+        res = collection.query(
+                    query_embeddings = [chunk],
+                    n_results=1,
+                    where={"key": {'$ne':main_key}},
+                    include=['metadatas','embeddings','distances']
                 )
+        
+        top_similar.append((main_text[idx],chunk,res,res['distances'][0]))
+    
+    most_similar_items = sorted(top_similar,key = lambda x:x[3])[:top_similar_count]
+
 
     logger.write("--------------- DONE -----------------\n")
     log_area.text(logger.getvalue())
-    return cosine_sim_scores
+
+    return most_similar_items
+
+
+
+
 
 # Streamlit Interface
 st.title("Check Infringement")
@@ -205,35 +244,95 @@ main_product = st.text_input('Enter Main Product Name', 'Philips led 7w bulb')
 main_url = st.text_input('Enter Main Product Manual URL', 'https://www.assets.signify.com/is/content/PhilipsConsumer/PDFDownloads/Colombia/technical-sheets/ODLI20180227_001-UPD-es_CO-Ficha_Tecnica_LED_MR16_Master_7W_Dim_12V_CRI90.pdf')
 search_method = st.selectbox('Choose Search Engine', ['All','duckduckgo', 'google', 'archive', 'github', 'wikipedia'])
 
-col1, col2 = st.columns(2)
+col1, col2, col3= st.columns(3)
 with col1:
     product_count = st.number_input("Number of Simliar Products",min_value=1, step=1, format="%i")
 with col2:
     link_count = st.number_input("Number of Links per product",min_value=1, step=1, format="%i")
+with col3:
+    need_image = st.selectbox("Process Images", ['True','False'])
 
-
-tag_option = st.selectbox('Choose Similarity Method', ["Complete Document Similarity","Field Wise Document Similarity"])
+top_similar_count = st.number_input("Top Similarities to be displayed",value=3,min_value=1, step=1, format="%i")
+tag_option = "Complete Document Similarity"
 
 
 if st.button('Check for Infringement'):
-    log_output = st.empty()  # Placeholder for log output
+    global log_output  # Placeholder for log output
+
+    tab1, tab2 = st.tabs(["Output", "Console"])
+
+    with tab2:
+        log_output = st.empty()
+
+    with tab1:
+        with st.spinner('Processing...'):
+            with StreamCapture() as logger:
+                top_similar_values = score(main_product, main_url, product_count, link_count, search_method, logger, log_output)
+
+        st.success('Processing complete!')
+
+        st.subheader("Cosine Similarity Scores")
+
+        for main_text, main_vector, response, _ in top_similar_values:
+            product_name = response['metadatas'][0][0]['product_name']
+            link = response['metadatas'][0][0]['url']
+            similar_text = response['metadatas'][0][0]['text']
+
+            cosine_score = cosine_similarity([main_vector], response['embeddings'][0])[0][0]
+
+            # Display the product information
+            with st.container():
+                st.markdown(f"### [Product: {product_name}]({link})")
+                st.markdown(f"#### Cosine Score: {cosine_score:.4f}")
+                col1, col2 = st.columns(2)
+                with col1:
+                    st.markdown(f"**Main Text:** {imporve_text(main_text)}")
+                with col2:
+                    st.markdown(f"**Similar Text:** {imporve_text(similar_text)}")
+
+                st.markdown("---")
+
+    if need_image == 'True':
+        with st.spinner('Processing Images...'):
+            emb_main = get_image_embeddings(main_product)
+            similar_prod = extract_similar_products(main_product)[0]
+            emb_similar = get_image_embeddings(similar_prod)
+
+            similarity_matrix = np.zeros((5, 5))
+            for i in range(5):
+                for j in range(5):
+                    similarity_matrix[i][j] = cosine_similarity([emb_main[i]], [emb_similar[j]])[0][0]
+
+            st.subheader("Image Similarity")
+            # Create an interactive heatmap
+            fig = px.imshow(similarity_matrix,
+                            labels=dict(x=f"{similar_prod} Images", y=f"{main_product} Images", color="Similarity"),
+                            x=[f"Image {i+1}" for i in range(5)],
+                            y=[f"Image {i+1}" for i in range(5)],
+                            color_continuous_scale="Viridis")
+
+            # Add title to the heatmap
+            fig.update_layout(title="Image Similarity Heatmap")
+
+            # Display the interactive heatmap
+            st.plotly_chart(fig)
+
+
+
+
+# main_product = 'Philips led 7w bulb'
+# main_url = 'https://www.assets.signify.com/is/content/PhilipsConsumer/PDFDownloads/Colombia/technical-sheets/ODLI20180227_001-UPD-es_CO-Ficha_Tecnica_LED_MR16_Master_7W_Dim_12V_CRI90.pdf'
+# search_method = 'duckduckgo'
+
+# product_count = 1
+# link_count = 1
+# need_image = False
+
 
-    with st.spinner('Processing...'):
-        with StreamCapture() as logger:
-            cosine_sim_scores = score(main_product, main_url,product_count, link_count, search_method, logger, log_output)
+# tag_option = "Field Wise Document Similarity"
 
-    st.success('Processing complete!')
+# logger = StreamCapture()
+# score(main_product, main_url,product_count, link_count, search_method, logger, st.empty())
 
-    st.subheader("Cosine Similarity Scores")
 
-    #  = score(main_product, main_url, search, logger, log_output)
-    if tag_option == 'Complete Document Similarity':
-        tags = ['Details']
-    else:    
-        tags = ['Introduction', 'Specifications', 'Product Overview', 'Safety Information', 'Installation Instructions', 'Setup and Configuration', 'Operation Instructions', 'Maintenance and Care', 'Troubleshooting', 'Warranty Information', 'Legal Information']
 
-    for product, link, index, value in cosine_sim_scores:
-        if not index:
-            st.write(f"Product: {product}, Link: {link}")
-        if value!=None:
-            st.write(f"{tags[index]:<20} - Similarity: {value:.2f}")