Spaces:
Runtime error
Runtime error
| 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 | |