Edit entity extraction modules

Author: JudeLee19

interact.py

@@ -93,7 +93,7 @@ def interact(self):
     def post_process(self, prediction, u_ent_features):
         if prediction == 0:
             return True
-        attr_list = [9, 11, 6, 1]
+        attr_list = [9, 12, 6, 1]
         if all(u_ent_featur == 1 for u_ent_featur in u_ent_features) and prediction in attr_list:
             return True
         else:
@@ -102,7 +102,7 @@ def post_process(self, prediction, u_ent_features):
     def action_post_process(self, prediction, u_entities):
         attr_mapping_dict = {
             9: '<cuisine>',
-            11: '<location>',
+            12: '<location>',
             6: '<party_size>',
             1: '<rest_type>'
         }

modules/actions.py

@@ -3,22 +3,22 @@
 
 '''
     Action Templates
-    'api_call <cuisine> <location> <party_size> <rest_type>',
-     '가격의 범위는 어느정도로 생각하세요',
-     '감사합니다',
-     '네 또 변경하실게 있나요',
-     '다른 리스트를 보여드릴게요',
-     '또 도와드릴게 있나요',
-     '몇명의 인원으로 예약하실 건가요',
-     '안녕하세요 어떻게 도와드릴까요',
-     '알겠습니다',
-     '어떤 종류의 요리를 좋아하나요',
-     '예약을 진행해드리도록 하겠습니다',
-     '위치는 <info_address> 입니다',
-     '위치는 어디에 있어야 하나요',
-     '이 리스트는 어떤가요: <restaurant>',
-     '전화번호는 <info_phone> 입니다',
-     '좋아요 몇 가지 리스트를 보여드릴게요'
+     0. api_call <cuisine> <location> <party_size> <rest_type>',
+     1. 가격의 범위는 어느정도로 생각하세요',
+     2. 감사합니다',
+     3. 네 또 변경하실게 있나요',
+     4. 다른 리스트를 보여드릴게요',
+     5. 또 도와드릴게 있나요',
+     6. 몇명의 인원으로 예약하실 건가요',
+     7. 안녕하세요 어떻게 도와드릴까요',
+     8. 알겠습니다',
+     9. 어떤 종류의 요리를 좋아하나요',
+     10. 예약을 진행해드리도록 하겠습니다',
+     11. 위치는 <info_address> 입니다',
+     12. 위치는 어디에 있어야 하나요',
+     13. 이 리스트는 어떤가요: <restaurant>',
+     14. 전화번호는 <info_phone> 입니다',
+     15. 좋아요 몇 가지 리스트를 보여드릴게요'
 '''
 
 

modules/entities.py

@@ -17,7 +17,7 @@ def __init__(self):
         # constants
         self.party_sizes = ['한명', '두명', '세명', '셋', '네명', '넷', '다섯', '다섯명', '여섯', '여섯명', '일곱', '여덟']
 
-        self.locations = ['방콕', '베이징', '붐베이', '하노이', '파리', '로마', '런던', '마드리드', '서울', '도쿄']
+        self.locations = ['방콕', '베이징', '붐베이', '하노이', '파리', '로마', '런던', '마드리드', '서울', '도쿄', 'LA']
 
         self.cuisines = ['영국','중국','프랑스', '이탈리아', '인도', '일식', '일본', '한식', '한국', '스페인', '타이', '베트남']
 
@@ -26,23 +26,28 @@ def __init__(self):
         self.EntType = Enum('Entity Type', '<party_size> <location> <cuisine> <rest_type> <non_ent>')
 
     def ent_type(self, ent):
+        # entity = [word for word in locations if word in input_str]
         if ent.startswith(tuple(self.party_sizes)):
-            return self.EntType['<party_size>'].name
+            entity_word = [word for word in self.party_sizes if word in ent][0]
+            return self.EntType['<party_size>'].name, entity_word
         elif ent.startswith(tuple(self.locations)):
-            return self.EntType['<location>'].name
+            entity_word = [word for word in self.locations if word in ent][0]
+            return self.EntType['<location>'].name, entity_word
         elif ent.startswith(tuple(self.cuisines)):
-            return self.EntType['<cuisine>'].name
+            entity_word = [word for word in self.cuisines if word in ent][0]
+            return self.EntType['<cuisine>'].name, entity_word
         elif ent.startswith(tuple(self.rest_types)):
-            return self.EntType['<rest_type>'].name
+            entity_word = [word for word in self.rest_types if word in ent][0]
+            return self.EntType['<rest_type>'].name, entity_word
         else:
-            return ent
+            return ent, None
 
     def extract_entities(self, utterance, update=True, is_test=False):
         tokenized = []
         for word in utterance.split(' '):
-            entity = self.ent_type(word)
+            entity, entity_word = self.ent_type(word)
             if word != entity and update:
-                self.entities[entity] = word
+                self.entities[entity] = entity_word
             tokenized.append(entity)
         tokenized_str = ' '.join(tokenized)