Feature branch to add feature for asynchronous inference with cv and llm models, for improving latency.

1. Added logic to separate main file from inference file, redirection in streamlit added. ( main_page and model_inference )
2. Moved loading of models to model_inference function.
3. Added asynchronous code for most of the inference and pre-processing functions to decrease latency.
4. Added threadpool executor to distribute llm inference, this saves about 25% latency time with llm.
5. Added logic to display inference of certain items that are done quickly than others, improves user experince, they need not wait for all items to be done to see result.
6. Added stateful application with page toggling.

Signed-off-by: Amith Adiraju <amithadiraju@Amiths-Laptop.local>

Documented most of the functions and cleaned up code.

Signed-off-by: Amith Adiraju <amithadiraju@Amiths-Laptop.local>

.gitignore

@@ -1,3 +1,4 @@
 misc.txt
 test_cas.py
-test_train_llm.py
+test_train_llm.py
+redir_app.py

app.py

@@ -1,62 +1,166 @@
 import streamlit as st
+from streamlit import session_state as sst
+from typing import List, Optional
+import asyncio
+import pandas as pd
 
 from inference.translate import (
     extract_filter_img,
-    transcribe_menu_model,
-    load_models
+    transcribe_menu_model
 )
 
 from inference.config import DEBUG_MODE
 from PIL import Image
 import time
 
-# Streamlit app
-st.title("Image Upload and Processing")
+from concurrent.futures import ThreadPoolExecutor, as_completed
+import os
 
+# Setting workers to be 70% of all available virtual cpus in system
+cpu_count = os.cpu_count()
+pool = ThreadPoolExecutor(max_workers=int(cpu_count*0.7) )
 
-# Using open source text detector, LLM for explaining items
-text_extractor, \
-    item_tokenizer,item_summarizer = load_models(item_summarizer = "google/flan-t5-large")
+# Initialize session state variable to start with home page
+if "page" not in sst:
+    sst["page"] = "Home"
 
-# Streamlit function to upload an image from any device
-uploaded_file = st.file_uploader("Choose an image...",
-                                 type=["jpg", "jpeg", "png"])
+def navigate_to(page: str) -> None:
+    """
+    Function to set the current page in the state of streamlit. A helper for 
+    simulating navigation in streamlit.
+
+    Parameters:
+        page: str, required.
+    
+    Returns:
+        None
+    """
+    
+    sst["page"] = page
+
+async def main_page() -> None:
+    """
+    Function that contains content of main page i.e., image uploader and submit button to navigate to next page.
+    Upon submit , control goes to model inference 'page'.
 
+    Parameters:
+        None
+    
+    Returns:
+        None
+    """
 
-# Submit button
-if uploaded_file is not None:
-    image = Image.open(uploaded_file)
+    # Streamlit app
+    first_title = st.empty()
+    first_title.title("App that explains your menu items ")
 
-    # Only show if user wants to see
-    if st.checkbox('Show Uploaded Image'):
-        st.image(image,
-                 caption='Uploaded Image',
-                 use_column_width=True)
+
+    # Streamlit function to upload an image from any device
+    uploaded_file = st.file_uploader("Choose an image...",
+                                 type=["jpg", "jpeg", "png"])
+    
+    # Remove preivous states' value of input image if it exists
+    sst.pop('input_image', None)
 
     # Submit button
-    if st.button("Submit"):
+    if uploaded_file is not None:
+        image = Image.open(uploaded_file)
+
+        # Only show if user wants to see
+        if st.checkbox('Show Uploaded Image'):
+            st.image(image,
+                    caption='Uploaded Image',
+                    use_column_width=True)
+
+        sst["input_image"] = image
+        
+        # Submit button
+        st.button("Submit",
+                  on_click = navigate_to,
+                  args = ("Inference",))
+        
+
+        st.info("""This application is for education purposes only. It uses AI, hence it's dietary
+                    recommendations are not to be taken as medical advice, author doesn't bear responsibility
+                    for incorrect dietary recommendations. Please proceed with caution.
+                    """)
+    
+
+async def dist_llm_inference(inp_texts: List[str]) -> None:
+
+    """ 
+    Function that performs concurrent LLM inference using threadpool. It displays 
+    results of those threads that are done with execution, as a dynamic row to streamlit table, rather than 
+    waiting for all threads to be done.
+
+    Parameters:
+        inp_texts: List[str], required -> List of strings, containing item names of a menu in english.
+
+    Returns:
+        None
+    """
+    
+    df = pd.DataFrame([('ITEM NAME', 'EXPLANATION')]
+                     )
+    
+    sl_table = st.table(df)
+    tp_futures = { pool.submit(transcribe_menu_model, mi): mi for mi in inp_texts }
+    
+    for tpftr in as_completed(tp_futures):
+
+        item = tp_futures[tpftr]
+
+        try:
+            exp = tpftr.result()
+            sl_table.add_rows([(item,exp)] )
+
+        except Exception as e:
+            print("Could not add a new row dynamically, because of this error:", e)
+    
+    return
+
+
+async def model_inference():
+    
+    """
+    Function that pre-processes input text from state variables, does concurrent inference
+    and toggles state between pages if needed.
+
+    Parameters:
+        None
+    Returns:
+        None
+
+    """
+    
+    second_title = st.empty()
+    second_title.title(" Using ML to explain your menu items ... ")
+    
+    if "input_image" in sst:
+        
+        image = sst["input_image"]
 
         msg1 = st.empty()
         msg1.write("Pre-processing and extracting text out of your image ....")
         st_filter = time.perf_counter()
+        
         # Call the extract_filter_img function
-        filtered_text = extract_filter_img(image, text_extractor)
+        filtered_text = await extract_filter_img(image)
         en_filter = time.perf_counter()
 
         num_items_detected = len(filtered_text)
+        
         if num_items_detected == 0:
             st.write("We couldn't detect any menu items ( indian for now ) from your image, please try a different image.")
-        
+
         elif num_items_detected > 0:
-            st.write(f"Detected {num_items_detected} menu items ( indian ) from your input image ... ")
+            st.write(f"Detected {num_items_detected} menu items from your input image ... ")
             
             msg2 = st.empty()
             msg2.write("All pre-processing done, transcribing your menu items now ....")
             st_trans_llm = time.perf_counter()
-            translated_text_dict = transcribe_menu_model(menu_texts=filtered_text,
-                                                        text_tokenizer=item_tokenizer,
-                                                        text_summarizer=item_summarizer
-                                                        )
+
+            await dist_llm_inference(filtered_text)
             
             msg3 = st.empty()                                            
             msg3.write("Done transcribing ... ")
@@ -74,5 +178,27 @@ if uploaded_file is not None:
                 st.write("Time took to summarize by LLM {}".format(llm_time_sec))
                 st.write('Overall time taken in seconds: {}'.format(total_time_sec))
             
-            st.table(translated_text_dict)
-        
+            
+            st.button("translate another",
+                      on_click=navigate_to,
+                      args=("Home",))        
+
+    else:
+        st.write("Looks like image upload failed, please try uploading it again ... ")
+
+
+async def main():
+    """
+    Function that toggles between pages based on state variables.
+
+    Parameters:
+        None
+    Returns:
+        None
+    """
+    if sst["page"] == "Home":
+        await main_page()
+    elif sst["page"] == "Inference":
+        await model_inference()
+
+asyncio.run(main())

inference/config.py

@@ -29,6 +29,5 @@ Based on Item and explanation pairs provided above, provide similar explanation
 Item ->
 """
 
-DEBUG_MODE = True
-
+DEBUG_MODE = False
 DEVICE = 'cpu'

inference/preprocess_image.py

@@ -11,6 +11,18 @@ import re
 
 
 def preprocess_text(sentence: AnyStr) -> AnyStr:
+
+    """ 
+    Function that pre-processes input text by removing special characters, hyper links,
+    numbers and by removing stop words
+
+    Parameters:
+        sentence: str, required -> A raw string which may have stop words, special chars etc.
+
+    Returns:
+        return_txt: str -> A clean string with all aforementioned, removed.
+    """
+
     sentence=sentence.lower().replace('{html}',"") 
     cleanr = re.compile('<.*?>')
     cleantext = re.sub(cleanr, '', sentence)
@@ -27,15 +39,25 @@ def preprocess_text(sentence: AnyStr) -> AnyStr:
     return return_txt
 
 def image_to_np_arr(image) -> np.array:
+
+    """ 
+    Function that converts a byte array image into a floating pointer numpy array.
+
+    Parameters:
+        inp_texts: List[str], required -> List of strings, containing item names of a menu in english.
+
+    Returns:
+        np.ndarray
+    """
+
     return np.array(image)
 
-def process_extracted_text(raw_extrc_text: List[Tuple]) -> List[AnyStr]:
+async def process_extracted_text(raw_extrc_text: List[Tuple]) -> List[AnyStr]:
     
     output_texts = []
     for _, extr_text, _ in raw_extrc_text:
         # remove all numbers, special characters from a string
         prcsd_txt = preprocess_text(extr_text)
-
-        if len(prcsd_txt.split(" ") ) > 2: output_texts.append(prcsd_txt)
+        if len(prcsd_txt.split(" ") ) >= 2: output_texts.append(prcsd_txt)
     
     return output_texts

inference/translate.py

@@ -14,9 +14,21 @@ import time
 use_gpu = True
 if DEVICE == 'cpu': use_gpu = False
 
+@st.cache_resource
+def load_models(item_summarizer: AnyStr) -> Tuple:
+    text_extractor = easyocr.Reader(['en'],
+                                    gpu = use_gpu
+                                    )
+    tokenizer = T5Tokenizer.from_pretrained(item_summarizer)
+    model = T5ForConditionalGeneration.from_pretrained(item_summarizer)
+
+    return (text_extractor, tokenizer, model)
+
+text_extractor,item_tokenizer,item_summarizer = load_models(item_summarizer = "google/flan-t5-large")
+
 
 # Define your extract_filter_img function
-def extract_filter_img(image, text_extractor) -> Dict:
+async def extract_filter_img(image) -> Dict:
 
     """
     1. Convert Image to numpy array
@@ -48,7 +60,8 @@ def extract_filter_img(image, text_extractor) -> Dict:
         if i in ind_add_delays:
             time.sleep(0.5)
 
-        result = func(result)
+        if i == 2: result = await func(result)
+        else: result = func(result)
         
         status_message.write(end_message)
 
@@ -63,42 +76,22 @@ def extract_filter_img(image, text_extractor) -> Dict:
     return result
 
 
-def transcribe_menu_model(menu_texts: List[AnyStr],
-                          text_summarizer = None,
-                          text_tokenizer = None) -> Dict:
-
-    summarized_menu_items = {}
+def transcribe_menu_model(menu_text: List[AnyStr]) -> Dict:
 
-    for mi in menu_texts:
-        if not text_summarizer:
-            raise NotImplementedError(""" """)
-        
-        else:
-            prompt_item = INSTRUCTION_PROMPT + " " + mi + """
+    prompt_item = INSTRUCTION_PROMPT + " " + menu_text + """
 
 
 """
-            input_ids = text_tokenizer(prompt_item, return_tensors="pt").input_ids
-            
-            outputs = text_summarizer.generate(input_ids,
-                                               max_new_tokens = 512
-                                               )
-            
-            summarized_menu_items[mi] = text_tokenizer.decode(
-                outputs[0],
-                skip_special_tokens = True
-                )
+    input_ids = item_tokenizer(prompt_item, return_tensors="pt").input_ids
     
-    return summarized_menu_items
-
-def load_models(item_summarizer: AnyStr) -> Tuple:
-    text_extractor = easyocr.Reader(['en'],
-                                    gpu = use_gpu
-                                    )
-    tokenizer = T5Tokenizer.from_pretrained(item_summarizer)
-    model = T5ForConditionalGeneration.from_pretrained(item_summarizer)
-
-    return (text_extractor, tokenizer, model)
+    outputs = item_summarizer.generate(input_ids,
+                                        max_new_tokens = 512
+                                        )
+    
+    return item_tokenizer.decode(
+        outputs[0],
+        skip_special_tokens = True
+        )
 
 def classify_menu_text(extrc_str: List[AnyStr]) -> List[AnyStr]:
     return extrc_str