added order taking NER

README.md

@@ -1,6 +1,6 @@
 ---
 title: Falcon Barista
-emoji: 👀
+emoji: 🦅
 colorFrom: red
 colorTo: gray
 sdk: gradio

app.py

@@ -1,7 +1,9 @@
 import gradio as gr
+import pandas as pd
 from asr_openai import AutomaticSpeechRecognition
 from tts_elevenlabs import ElevenLabsTTS
 from falcon_7b_llm import Falcon_7b_llm
+from order_parser import Order_Parser
 import logging
 import os
 
@@ -21,7 +23,15 @@ def generate_response(input_audio):
     sentence = asr.run_transcription(input_audio)
     # sentence = 'how are you?'
     print(sentence)
+    global order_dict
+    try:
+        order_dict = order_taking.order_parser(sentence['text'])
+        print(order_dict)
+    except Exception as e:
+        print('order parsing failed')
+        print(e)
     llm_response = llm.get_llm_response(sentence['text'])
+    print(llm_response)
     output_audio = tts.tts_generate_audio(llm_response)
     # output_audio = tts.tts_generate_audio(sentence)
     chatbot_history.append(((input_audio,), (output_audio,)))
@@ -29,33 +39,71 @@ def generate_response(input_audio):
 
 delete_files_in_folder('data//tts_responses')
 
-title = "<h1 style='text-align: center; color: #ffffff; font-size: 40px;'> 🦅 Falcon Barista"
+title = "<h1 style='text-align: center; color: #ffffff; font-size: 40px;'> Falcon Barista (Pre-Alpha Release)"
 
 asr = AutomaticSpeechRecognition()
 tts = ElevenLabsTTS()
 llm = Falcon_7b_llm()
+order_taking = Order_Parser()
 chatbot_history = []
+order_display=False
+order_dict={}
 
-def restart_chat():
-    delete_files_in_folder('data//tts_responses')
-    global chatbot_history
-    chatbot_history = []
-    tts.restart_state()
-    llm.restart_state()
-    return chatbot_history
+df = pd.DataFrame({
+    "item" : [], 
+    "quantity" : [], 
+}) 
+
+s = df#.style.format("{:.2f}")
 
 with gr.Blocks() as demo:
     gr.Markdown(title)
+    order_title = gr.Markdown('### Your Order', visible=False)
     with gr.Row():
         gr.Image('https://i.imgur.com/fHCFI2T.png', label="Look how cute is Falcon Barista")
-        # gr.Image('data//falcon.png', label="Look how cute is Falcon Barista")
         with gr.Column():
-            chatbot = gr.Chatbot(label='Chat with Falcon Barista', avatar_images=('data//user_avatar_logo.png','data//falcon_logo_transparent.png'))
-            with gr.Row():
-                mic = gr.Audio(source="microphone", type='filepath', scale=3)
-                mic.stop_recording(generate_response, mic, chatbot)
-                restart_btn = gr.Button(value="Restart Chat", scale=1)
-                restart_btn.click(restart_chat, outputs=[chatbot])
+            chatbot = gr.Chatbot(label='Chat with Falcon Barista', avatar_images=('data//user_avatar_logo.png','data//falcon_logo_transparent.png'), scale=2)
+            mic = gr.Audio(source="microphone", type='filepath', scale=1)
+            mic.stop_recording(generate_response, mic, chatbot)
+    with gr.Row():
+        restart_btn = gr.Button(value="Restart Chat", scale=1, variant='stop')
+        # restart_btn.click(restart_chat, outputs=[chatbot])
+        end_btn = gr.Button(value="End Chat and Confirm Order", scale=2, variant='primary')
+        
+    with gr.Column(visible=False) as output_col:
+        order_title = gr.Markdown('### Your Order')
+        order_summary = gr.DataFrame(s)
+
+    def restart_chat():
+        delete_files_in_folder('data//tts_responses')
+        global chatbot_history
+        chatbot_history = []
+        global order_dict
+        order_dict = {}
+        global df
+        df = pd.DataFrame({
+            "item" : [], 
+            "quantity" : [], 
+        })
+        order_taking.restart_state() 
+        tts.restart_state()
+        llm.restart_state()
+        return {
+            chatbot: [],
+            output_col: gr.Column(visible=False)
+        }
+
+    def end_chat():
+        df = pd.DataFrame(list(order_dict.items()), columns=['item', 'quantity'])
+
+        s = df#.style.format("{:.2f}")
+        return {
+            output_col: gr.Column(visible=True),
+            order_summary: gr.DataFrame(s, visible=True)
+            }
+
+    restart_btn.click(restart_chat, outputs=[chatbot, output_col])
+    end_btn.click(end_chat, outputs=[output_col, order_summary])
 
 if __name__ == "__main__":
     demo.launch()

order_parser.py

@@ -0,0 +1,87 @@
+from transformers import AutoTokenizer, AutoModelForTokenClassification
+from transformers import pipeline
+from word2number import w2n
+import pandas as pd
+
+class Order_Parser():
+    def __init__(self):
+        tokenizer = AutoTokenizer.from_pretrained("davanstrien/deberta-v3-base_fine_tuned_food_ner")
+        model = AutoModelForTokenClassification.from_pretrained("davanstrien/deberta-v3-base_fine_tuned_food_ner")
+        self.pipe = pipeline("ner", model=model, tokenizer=tokenizer)
+        self.complete_order_dict={}
+
+    def restart_state(self):
+        self.complete_order_dict={}
+
+    def join_adjacent_items(self, data):
+        result = []
+        current_group = []
+        current_entity = None
+
+        for item in data:
+            # Check if the item's entity is related to FOOD or QUANTITY
+            if any(e in item['entity'] for e in ['FOOD', 'QUANTITY']):
+                # Start a new group if the entity type changes
+                if not current_entity:
+                    current_entity = item['entity'].split('-')[-1]
+                elif current_entity != item['entity'].split('-')[-1]:
+                    result.append({'entity': current_entity, 'word': ''.join(current_group)})
+                    current_group = []
+                    current_entity = item['entity'].split('-')[-1]
+                    
+                current_group.append(item['word'])
+            else:
+                if current_group:
+                    result.append({'entity': current_entity, 'word': ''.join(current_group)})
+                    current_group = []
+                    current_entity = None
+                result.append(item)
+
+        # Handle the last group if it exists
+        if current_group:
+            result.append({'entity': current_entity, 'word': ''.join(current_group)})
+
+        return result
+    
+    def order_parser(self, sentence):
+        sentence = sentence.replace(',', ' ')
+        sentence = sentence.replace('?', ' ')
+        sentence = sentence.replace('.', ' ')
+        # updated_sentence = updated_sentence.replace('and', 'one')
+        sentence = sentence.replace(' a ', ' one ')
+        sentence = sentence.replace(' an ', ' one ')
+        # sentence = sentence.replace(' and ', ' one ')
+        # print(updated_sentence)
+        # raw_order = self.pipe(updated_sentence)
+        print(sentence)
+        raw_order = self.pipe(sentence)
+        raw_order_piped = self.join_adjacent_items(raw_order)
+        
+        order_dict={}
+        quantity_exist = False
+        for ent in raw_order_piped:
+            if 'QUANTITY' in ent['entity']:
+                quantity = ent['word'].replace('▁', ' ')
+                try:
+                    quantity = w2n.word_to_num(quantity)
+                except:
+                    quantity = 1
+
+                # print(quantity)
+                quantity_exist = True
+
+            elif 'FOOD' in ent['entity']:
+                food = ent['word'].replace('▁', '')
+                # print(food)
+                
+                if quantity_exist:
+                    order_dict[food] = quantity
+                else:
+                    order_dict[food] = 1
+
+                quantity_exist=False
+
+        # print(order_dict)
+        self.complete_order_dict.update(order_dict)
+        return self.complete_order_dict
+    

requirements.txt

@@ -8,4 +8,5 @@ scipy
 elevenlabs
 openai
 torch
-wandb
+wandb
+word2number