Upload app.py

app.py

@@ -0,0 +1,234 @@
+
+import pandas as pd
+import numpy as np
+import matplotlib.pyplot as plt
+import seaborn as sns
+import plotly.express as px
+
+num_rows = 20000  
+df = pd.read_csv('/emails.csv', on_bad_lines='skip', nrows=num_rows)
+
+def get_message(Series: pd.Series):
+    result = pd.Series(index=Series.index)
+    for row, message in enumerate(Series):
+        message_words = message.split('\n')
+        del message_words[:15]
+        result.iloc[row] = ''.join(message_words).strip()
+    return result
+
+def get_date(Series: pd.Series):
+    result = pd.Series(index=Series.index)
+    for row, message in enumerate(Series):
+        message_words = message.split('\n')
+        del message_words[0]
+        del message_words[1:]
+        result.iloc[row] = ''.join(message_words).strip()
+        result.iloc[row] = result.iloc[row].replace('Date: ', '')
+    print('Done parsing, converting to datetime format..')
+    return pd.to_datetime(result)
+
+def get_sender_and_receiver(Series: pd.Series):
+    sender = pd.Series(index = Series.index)
+    recipient1 = pd.Series(index = Series.index)
+    recipient2 = pd.Series(index = Series.index)
+    recipient3 = pd.Series(index = Series.index)
+
+    for row,message in enumerate(Series):
+        message_words = message.split('\n')
+        sender[row] = message_words[2].replace('From: ', '')
+        recipient1[row] = message_words[3].replace('To: ', '')
+        recipient2[row] = message_words[10].replace('X-cc: ', '')
+        recipient3[row] = message_words[11].replace('X-bcc: ', '')
+
+    return sender, recipient1, recipient2, recipient3
+
+def get_subject(Series: pd.Series):
+    result = pd.Series(index = Series.index)
+
+    for row, message in enumerate(Series):
+        message_words = message.split('\n')
+        message_words = message_words[4]
+        result[row] = message_words.replace('Subject: ', '')
+    return result
+
+def get_folder(Series: pd.Series):
+    result = pd.Series(index = Series.index)
+
+    for row, message in enumerate(Series):
+        message_words = message.split('\n')
+        message_words = message_words[12]
+        result[row] = message_words.replace('X-Folder: ', '')
+    return result
+
+df['text'] = get_message(df.message)
+df['sender'], df['recipient1'], df['recipient2'], df['recipient3'] = get_sender_and_receiver(df.message)
+df['Subject'] = get_subject(df.message)
+df['folder'] = get_folder(df.message)
+df['date'] = get_date(df.message)
+
+df = df.drop(['message', 'file'], axis = 1)
+
+df.head(100)
+
+
+import chromadb
+chroma_client = chromadb.Client()
+
+collection = chroma_client.create_collection(name="emails")
+
+df.loc[4, 'text']
+
+for i in df.index:
+    collection.add(
+     
+        documents = df.loc[i, 'text'],
+
+      
+        metadatas = [{"sender": df.loc[i, 'sender'],
+                     "recipient1": df.loc[i, 'recipient1'],
+                     "recipient2": df.loc[i, 'recipient2'],
+                     "recipient3": df.loc[i, 'recipient3'],
+                     "subject": df.loc[i, 'Subject'],
+                     "folder": df.loc[i, 'folder'],
+                     "date": str(df.loc[i, 'date'])
+                     }],
+
+     
+        ids = str(i)
+    )
+
+collection.get(
+    ids=["140"]
+)
+
+results = collection.query(
+    query_texts = ["this is a document"],
+    n_results = 2,
+    include = ['distances', 'metadatas', 'documents']
+)
+results
+
+
+from chromadb.utils import embedding_functions
+
+
+
+sentence_transformer_ef = embedding_functions.SentenceTransformerEmbeddingFunction(model_name="paraphrase-MiniLM-L3-v2")
+
+collection_minilm = chroma_client.create_collection(name="emails_minilm", embedding_function=sentence_transformer_ef)
+
+
+for i in df.index:
+    print(i)
+    collection_minilm.add(
+       
+        documents = df.loc[i, 'text'],
+
+        metadatas = [{"sender": df.loc[i, 'sender'],
+                     "recipient1": df.loc[i, 'recipient1'],
+                     "recipient2": df.loc[i, 'recipient2'],
+                     "recipient3": df.loc[i, 'recipient3'],
+                     "subject": df.loc[i, 'Subject'],
+                     "folder": df.loc[i, 'folder'],
+                     "date": str(df.loc[i, 'date'])
+                     }],
+
+    
+        ids = str(i)
+    )
+
+results = collection_minilm.query(
+    query_texts = ["this is a document"],
+    n_results = 2,
+    include = ['distances', 'metadatas', 'documents']
+)
+results
+
+
+
+
+
+import gradio as gr
+
+
+def query_chromadb(question,numberOfResults):
+    results = collection_minilm.query(
+          n_results = numberOfResults,
+      )
+
+    return results['documents'][0]
+
+iface = gr.Interface(
+    fn=query_chromadb,
+    inputs=["text","number"],
+    outputs="text",
+    title="Email Dataset Interface",
+    description="Insert the question or the key word to find the topic correlated in the dataset"
+)
+
+iface.launch(share=True)
+
+
+
+import ast
+
+def create_output(dictionary, number):
+
+    dictionary_ids = str(dictionary['ids'])
+
+  
+    dictionary_ids_clean = dictionary_ids.strip("[]")
+
+    dictionary_ids_clean = dictionary_ids_clean.replace("'", "")
+
+   
+    dictionary_ids_list = dictionary_ids_clean.split(", ")
+
+    string_results = "";
+
+
+    for n in range(number):
+      t = collection_minilm.get(
+              ids=[dictionary_ids_list[n]]
+          )
+
+
+      id = str(t["ids"])
+      doc = str(t["documents"])
+      metadata = str(t["metadatas"])
+
+      dictionary_metadata = ast.literal_eval(metadata.strip("[]"))
+
+      string_results_old = string_results
+
+      string_temp = """---------------
+      SUBJECT: """ + dictionary_metadata['subject'] + """"
+      MESSAGE: """ + "\n" + doc + """
+      ---------------"""
+
+      string_results = string_results_old + string_temp
+
+    return string_results
+
+def query_chromadb_advanced(question,numberOfResults):
+    results = collection_minilm.query(
+          query_texts = question,
+          n_results = numberOfResults,
+      )
+
+    return create_output(results, numberOfResults)
+
+
+result_advance = query_chromadb_advanced("bank", 4)
+
+print(result_advance)
+
+iface = gr.Interface(
+    fn=query_chromadb_advanced,
+    inputs=["text","number"],
+    outputs="text",
+    title="Email Dataset Interface",
+    description="Insert the question or the key word to find the topic correlated in the dataset"
+)
+
+iface.launch(share=True)