Adds file support (txt, pdf, csv)

interface/components.py

@@ -1,5 +1,5 @@
 import streamlit as st
-from interface.utils import get_pipelines, extract_text_from_url
+from interface.utils import get_pipelines, extract_text_from_url, extract_text_from_file
 from interface.draw_pipelines import get_pipeline_graph
 
 
@@ -84,3 +84,27 @@ def component_article_url(container):
             {"text": doc["text"], "id": doc_id} for doc_id, doc in enumerate(urls)
         ]
         return corpus
+
+
+def component_file_input(container):
+    """Draw the extract text from file widget"""
+    with container:
+        files = []
+        doc_id = 1
+        with st.expander("Enter Files"):
+            while True:
+                file = st.file_uploader("Upload a .txt, .pdf, .csv file", key=doc_id)
+                if file != None:
+                    extracted_text = extract_text_from_file(file)
+                    if extracted_text != None:
+                        files.append({"text": extracted_text})
+                        doc_id += 1
+                        st.markdown("---")
+                    else:
+                        break
+                else:
+                    break
+        corpus = [
+            {"text": doc["text"], "id": doc_id} for doc_id, doc in enumerate(files)
+        ]
+        return corpus

interface/pages.py

@@ -2,6 +2,7 @@ import streamlit as st
 from streamlit_option_menu import option_menu
 from core.search_index import index, search
 from interface.components import (
+    component_file_input,
     component_show_pipeline,
     component_show_search_result,
     component_text_input,
@@ -59,6 +60,7 @@ def page_index(container):
         input_funcs = {
             "Raw Text": (component_text_input, "card-text"),
             "URL": (component_article_url, "card-link"),
+            "File": (component_file_input, "card-file"),
         }
         selected_input = option_menu(
             "Input Text",

interface/utils.py

@@ -1,7 +1,10 @@
+from io import StringIO
 import core.pipelines as pipelines_functions
 from inspect import getmembers, isfunction
 from newspaper import Article
+from PyPDF2 import PdfFileReader
 import streamlit as st
+import pandas as pd
 
 
 def get_pipelines():
@@ -21,3 +24,60 @@ def extract_text_from_url(url: str):
     article.parse()
 
     return article.text
+
+
+def extract_text_from_file(file):
+    # read text file
+    if file.type == "text/plain":
+        # To convert to a string based IO:
+        stringio = StringIO(file.getvalue().decode("utf-8"))
+
+        # To read file as string:
+        file_text = stringio.read()
+
+        return file_text
+
+    # read pdf file
+    elif file.type == "application/pdf":
+        pdfReader = PdfFileReader(file)
+        count = pdfReader.numPages
+        all_text = ""
+
+        for i in range(count):
+            try:
+                page = pdfReader.getPage(i)
+                all_text += page.extractText()
+            except:
+                continue
+        file_text = all_text
+
+        return file_text
+    
+    # read csv file
+    elif file.type == "text/csv":
+        csv = pd.read_csv(file)
+        # get columns of type string
+        string_columns = csv.select_dtypes(include=['object']).columns
+        # get data from columns and join it together
+        file_text = ""
+        for row in csv[string_columns].values.tolist():
+            # remove NaNs
+            row = [x for x in row if str(x) != 'nan']
+            for column in row:
+                txt = ""
+                if isinstance(column, list):
+                    try:
+                       txt = " ".join(column)
+                    except:
+                        continue
+                elif isinstance(column, str):
+                    txt = column
+                else:
+                    continue
+                file_text += " " + txt
+        return file_text
+
+    else:
+        st.warning(f"File type {file.type} not supported")
+        return None
+

requirements.txt

@@ -3,4 +3,5 @@ streamlit_option_menu==0.3.2
 farm-haystack==1.8.0
 black==22.8.0
 plotly==5.10.0
-newspaper3k==0.2.8
+newspaper3k==0.2.8
+PyPDF2==2.10.7