final touches to draw pipelines & manual cache

core/pipelines.py

@@ -2,15 +2,12 @@
 Haystack Pipelines
 """
 
-import tokenizers
 from haystack import Pipeline
 from haystack.document_stores import InMemoryDocumentStore
 from haystack.nodes.retriever import DensePassageRetriever, TfidfRetriever
 from haystack.nodes.preprocessor import PreProcessor
-import streamlit as st
 
 
-@st.cache(allow_output_mutation=True)
 def keyword_search(
     index="documents",
 ):
@@ -42,13 +39,6 @@ def keyword_search(
     return search_pipeline, index_pipeline
 
 
-@st.cache(
-    hash_funcs={
-        tokenizers.Tokenizer: lambda _: None,
-        tokenizers.AddedToken: lambda _: None,
-    },
-    allow_output_mutation=True,
-)
 def dense_passage_retrieval(
     index="documents",
     query_embedding_model="facebook/dpr-question_encoder-single-nq-base",

interface/components.py

@@ -13,10 +13,16 @@ def component_select_pipeline(container):
             if "Keyword Search" in pipeline_names
             else 0,
         )
-        (
-            st.session_state["search_pipeline"],
-            st.session_state["index_pipeline"],
-        ) = pipeline_funcs[pipeline_names.index(selected_pipeline)]()
+        if st.session_state["pipeline"] is None or st.session_state["pipeline"]["name"] != selected_pipeline:
+            (
+                search_pipeline,
+                index_pipeline,
+            ) = pipeline_funcs[pipeline_names.index(selected_pipeline)]()
+            st.session_state["pipeline"] = {
+                'name': selected_pipeline,
+                'search_pipeline': search_pipeline,
+                'index_pipeline': index_pipeline,
+            }
 
 
 def component_show_pipeline(pipeline):

interface/config.py

@@ -1,7 +1,9 @@
 from interface.pages import page_landing_page, page_search, page_index
 
 # Define default Session Variables over the whole session.
-session_state_variables = {}
+session_state_variables = {
+    "pipeline": None
+}
 
 # Define Pages for the demo
 pages = {

interface/draw_pipelines.py

@@ -3,11 +3,9 @@ from typing import List
 from itertools import chain
 import networkx as nx
 import plotly.graph_objs as go
-import streamlit as st
 import numpy as np
 
 
-@st.cache(allow_output_mutation=True)
 def get_pipeline_graph(pipeline):
     # Controls for how the graph is drawn
     nodeColor = "#ffbf00"
@@ -16,13 +14,37 @@ def get_pipeline_graph(pipeline):
     lineColor = "#ffffff"
 
     G = pipeline.graph
-    initial_coordinate = (0, len(G.nodes))
-    fixed_pos = {
-        node: np.array([initial_coordinate[0], initial_coordinate[1] - float(idx)])
-        for idx, node in enumerate(G.nodes)
-    }
-    pos = nx.spring_layout(G, pos=fixed_pos, seed=42)
-
+    current_coordinate = (0, len(set([edge[0] for edge in G.edges()])) + 1)
+    # Transform G.edges into {node : all_connected_nodes} format
+    node_connections = {}
+    for in_node, out_node in G.edges():
+        if in_node in node_connections:
+            node_connections[in_node].append(out_node)
+        else:
+            node_connections[in_node] = [out_node]
+    # Get node coordinates/pos
+    fixed_pos_nodes = {}
+    for idx, (in_node, out_nodes) in enumerate(node_connections.items()):
+        if in_node not in fixed_pos_nodes:
+            fixed_pos_nodes[in_node] = np.array([current_coordinate[0], current_coordinate[1]])
+            current_coordinate = (current_coordinate[0], current_coordinate[1] - 1)
+        # If more than 1 out node, then branch out in X coordinate
+        if len(out_nodes) > 1:
+            # if length is odd
+            if (len(out_nodes) % 2) != 0:
+                middle_node = out_nodes[round(len(out_nodes)/2, 0) - 1]
+                fixed_pos_nodes[middle_node] = np.array([current_coordinate[0], current_coordinate[1]])
+                out_nodes = [n for n in out_nodes if n != middle_node]
+            correction_coordinate = - len(out_nodes) / 2
+            for out_node in out_nodes:
+                fixed_pos_nodes[out_node] = np.array([int(current_coordinate[0] + correction_coordinate), int(current_coordinate[1])])
+                if correction_coordinate == -1:
+                    correction_coordinate += 1
+                correction_coordinate += 1
+            current_coordinate = (current_coordinate[0], current_coordinate[1] - 1)
+        elif len(node_connections) - 1 == idx:
+            fixed_pos_nodes[out_nodes[0]] = np.array([current_coordinate[0], current_coordinate[1]])
+    pos = nx.spring_layout(G, pos=fixed_pos_nodes, fixed=G.nodes(), seed=42)
     for node in G.nodes:
         G.nodes[node]["pos"] = list(pos[node])
 

interface/pages.py

@@ -36,12 +36,12 @@ def page_search(container):
         ## SEARCH ##
         query = st.text_input("Query")
 
-        component_show_pipeline(st.session_state["search_pipeline"])
+        component_show_pipeline(st.session_state["pipeline"]["search_pipeline"])
 
         if st.button("Search"):
             st.session_state["search_results"] = search(
                 queries=[query],
-                pipeline=st.session_state["search_pipeline"],
+                pipeline=st.session_state["pipeline"]["search_pipeline"],
             )
         if "search_results" in st.session_state:
             component_show_search_result(
@@ -53,7 +53,7 @@ def page_index(container):
     with container:
         st.title("Index time!")
 
-        component_show_pipeline(st.session_state["index_pipeline"])
+        component_show_pipeline(st.session_state["pipeline"]["index_pipeline"])
 
         input_funcs = {
             "Raw Text": (component_text_input, "card-text"),
@@ -74,7 +74,7 @@ def page_index(container):
             if st.button("Index"):
                 index_results = index(
                     corpus,
-                    st.session_state["index_pipeline"],
+                    st.session_state["pipeline"]["index_pipeline"],
                 )
             if index_results:
                 st.write(index_results)