Restore working version

app.py

@@ -1,13 +1,249 @@
 import gradio as gr
-from predict import run_episode
+import time, torch
+from transformers import BartTokenizer, BartForConditionalGeneration, AutoModel, AutoTokenizer
+
+from webshop_lite import dict_to_fake_html
+from predict_help import convert_dict_to_actions, convert_html_to_text, parse_results, parse_item_page, Page
+
+# load IL models
+bart_tokenizer = BartTokenizer.from_pretrained('facebook/bart-large')
+bart_model = BartForConditionalGeneration.from_pretrained('webshop/il_search_bart')
+
+bert_tokenizer = AutoTokenizer.from_pretrained('bert-base-uncased', truncation_side='left')
+bert_tokenizer.add_tokens(['[button]', '[button_]', '[clicked button]', '[clicked button_]'], special_tokens=True)
+bert_model = AutoModel.from_pretrained('webshop/il-choice-bert-image_0', trust_remote_code=True)
+
+def process_str(s):
+    s = s.lower().replace('"', '').replace("'", "").strip()
+    s = s.replace('[sep]', '[SEP]')
+    return s
+
+
+def process_goal(state):
+    state = state.lower().replace('"', '').replace("'", "")
+    state = state.replace('amazon shopping game\ninstruction:', '').replace('\n[button] search [button_]', '').strip()
+    if ', and price lower than' in state:
+        state = state.split(', and price lower than')[0]
+    return state
+
+
+def data_collator(batch):
+    state_input_ids, state_attention_mask, action_input_ids, action_attention_mask, sizes, labels, images = [], [], [], [], [], [], []
+    for sample in batch:
+        state_input_ids.append(sample['state_input_ids'])
+        state_attention_mask.append(sample['state_attention_mask'])
+        action_input_ids.extend(sample['action_input_ids'])
+        action_attention_mask.extend(sample['action_attention_mask'])
+        sizes.append(sample['sizes'])
+        labels.append(sample['labels'])
+        images.append(sample['images'])
+    max_state_len = max(sum(x) for x in state_attention_mask)
+    max_action_len = max(sum(x) for x in action_attention_mask)
+    return {
+        'state_input_ids': torch.tensor(state_input_ids)[:, :max_state_len],
+        'state_attention_mask': torch.tensor(state_attention_mask)[:, :max_state_len],
+        'action_input_ids': torch.tensor(action_input_ids)[:, :max_action_len],
+        'action_attention_mask': torch.tensor(action_attention_mask)[:, :max_action_len],
+        'sizes': torch.tensor(sizes),
+        'images': torch.tensor(images),
+        'labels': torch.tensor(labels),
+    }
+
+
+def bart_predict(input):
+    input_ids = bart_tokenizer(input)['input_ids']
+    input_ids = torch.tensor(input_ids).unsqueeze(0)
+    output = bart_model.generate(input_ids, max_length=512, num_return_sequences=5, num_beams=5)
+    return bart_tokenizer.batch_decode(output.tolist(), skip_special_tokens=True)[0]
+
+
+def bert_predict(obs, info, softmax=True):
+    valid_acts = info['valid']
+    assert valid_acts[0].startswith('click[')
+    state_encodings = bert_tokenizer(process_str(obs), max_length=512, truncation=True, padding='max_length')
+    action_encodings = bert_tokenizer(list(map(process_str, valid_acts)), max_length=512, truncation=True,  padding='max_length')
+    batch = {
+        'state_input_ids': state_encodings['input_ids'],
+        'state_attention_mask': state_encodings['attention_mask'],
+        'action_input_ids': action_encodings['input_ids'],
+        'action_attention_mask': action_encodings['attention_mask'],
+        'sizes': len(valid_acts),
+        'images': info['image_feat'].tolist(),
+        'labels': 0
+    }
+    batch = data_collator([batch])
+    outputs = bert_model(**batch)
+    if softmax:
+        idx = torch.multinomial(torch.nn.functional.softmax(outputs.logits[0], dim=0), 1)[0].item()
+    else:
+        idx = outputs.logits[0].argmax(0).item()
+    return valid_acts[idx]
+
+
+def predict(obs, info):
+    """
+    Given WebShop environment observation and info, predict an action.
+    """
+    valid_acts = info['valid']
+    if valid_acts[0].startswith('click['):
+        return bert_predict(obs, info)
+    else:
+        return "search[" + bart_predict(process_goal(obs)) + "]"
+
+def run_episode(goal, verbose=True):
+    """
+    Interact with amazon to find a product given input goal.
+    Input: text goal
+    Output: a url of found item on amazon.
+    """
+    obs = "Amazon Shopping Game\nInstruction:" + goal + "\n[button] search [button]"
+    info = {'valid': ['search[stuff]'], 'image_feat': torch.zeros(512)}
+    product_map = {}
+    title_to_asin_map = {}
+    search_results_cache = {}
+    visited_asins, clicked_options = set(), set()
+    sub_page_type, page_type, page_num = None, None, None
+    search_terms, prod_title, asin, num_prods, = None, None, None, None
+    options = {}
+    
+    for i in range(100):
+        # Run prediction
+        action = predict(obs, info)
+        if verbose:
+            print("====")
+            print(action)
+        
+        # Previous Page Type, Action -> Next Page Type
+        action_content = action[action.find("[")+1:action.find("]")]
+        prev_page_type = page_type
+        if action.startswith('search['):
+            page_type = Page.RESULTS
+            search_terms = action_content
+            page_num = 1
+        elif action.startswith('click['):
+            if action.startswith('click[item -'):
+                prod_title = action_content[len("item -"):].strip()
+                found = False
+                for key in title_to_asin_map:
+                    if prod_title == key:
+                        asin = title_to_asin_map[key]
+                        page_type = Page.ITEM_PAGE
+                        visited_asins.add(asin)
+                        found = True
+                        break
+                if not found:
+                    raise Exception("Product to click not found")
+                    
+            elif any(x.value in action for x in [Page.DESC, Page.FEATURES, Page.REVIEWS]):
+                page_type = Page.SUB_PAGE
+                sub_page_type = Page(action_content.lower())
+                
+            elif action == 'click[< prev]':
+                if sub_page_type is not None:
+                    page_type, sub_page_type = Page.ITEM_PAGE, None
+                elif prev_page_type == Page.ITEM_PAGE:
+                    page_type = Page.RESULTS
+                    options, clicked_options = {}, set()
+                elif prev_page_type == Page.RESULTS and page_num > 1:
+                    page_type = Page.RESULTS
+                    page_num -= 1
+                    
+            elif action == 'click[next >]':
+                page_type = Page.RESULTS
+                page_num += 1
+                
+            elif action.lower() == 'click[back to search]':
+                page_type = Page.SEARCH
+                
+            elif action == 'click[buy now]':
+                asin_url = f"https://www.amazon.com/dp/{asin}"
+                return_value = "Product URL: " + asin_url
+                if len(clicked_options) > 0:    
+                    options_str = ', '.join(list(clicked_options))
+                    return_value += "\nSelected Options: " + options_str
+                return return_value
+            
+            elif prev_page_type == Page.ITEM_PAGE:
+                found = False
+                for opt_name, opt_values in product_map[asin]["options"].items():
+                    if action_content in opt_values:
+                        options[opt_name] = action_content
+                        page_type = Page.ITEM_PAGE
+                        clicked_options.add(action_content)
+                        found = True
+                        break
+                if not found:
+                    raise Exception("Unrecognized action: " + action)
+        else:
+            raise Exception("Unrecognized action:" + action)
+        
+        if verbose:
+            print(f"Parsing {page_type.value} page...")
+        
+        # URL -> Real HTML -> Dict of Info
+        if page_type == Page.RESULTS:
+            if search_terms in search_results_cache:
+                data = search_results_cache[search_terms]
+            else:
+                begin = time.time()
+                data = parse_results(search_terms, page_num)
+                end = time.time()
+                print("Parsing search results took", end-begin, "seconds")
+
+                search_results_cache[search_terms] = data
+                num_prods = len(data)
+                for d in data:
+                    title_to_asin_map[d['Title']] = d['asin']
+        elif page_type == Page.ITEM_PAGE or page_type == Page.SUB_PAGE:
+            if asin in product_map:
+                print("Loading cached item page for", asin)
+                data = product_map[asin]
+            else:
+                begin = time.time()
+                data = parse_item_page(asin)
+                end = time.time()
+                print("Parsing item page took", end-begin, "seconds")
+                product_map[asin] = data
+        elif page_type == Page.SEARCH:
+            if verbose:
+                print("Executing search")
+            obs = "Amazon Shopping Game\nInstruction:" + goal + "\n[button] search [button]"
+            info = {'valid': ['search[stuff]'], 'image_feat': torch.zeros(512)}
+            continue
+        else:
+            raise Exception("Page of type `", page_type, "` not found")
+
+        # Dict of Info -> Fake HTML -> Text Observation
+        begin = time.time()
+        html_str = dict_to_fake_html(data, page_type, asin, sub_page_type, options, product_map, goal)
+        obs = convert_html_to_text(html_str, simple=False, clicked_options=clicked_options, visited_asins=visited_asins)
+        end = time.time()
+        print("[Page Info -> WebShop HTML -> Observation] took", end-begin, "seconds")
+
+        # Dict of Info -> Valid Action State (Info)
+        begin = time.time()
+        prod_arg = product_map if page_type == Page.ITEM_PAGE else data
+        info = convert_dict_to_actions(page_type, prod_arg, asin, page_num, num_prods)
+        end = time.time()
+        print("Extracting available actions took", end-begin, "seconds")
+        
+        if i == 99:
+            asin_url = f"https://www.amazon.com/dp/{asin}"
+            return_value = "Product URL: " + asin_url
+            if len(clicked_options) > 0:    
+                options_str = ', '.join(list(clicked_options))
+                return_value += "\nSelected Options: " + options_str
+            return return_value
 
 gr.Interface(fn=run_episode,\
     inputs=gr.inputs.Textbox(lines=7, label="Input Text"),\
     outputs="text",\
     examples=[
+        "Please select a 1 pound, certified organic sea salt shaker in the flavor triple blend flakes, and price lower than 40.00 dollars",
         "I want to find a gold floor lamp with a glass shade and a nickel finish that i can use for my living room, and price lower than 270.00 dollars",
         "I'm trying to find white bluetooth speakers that are not only water resistant but also come with stereo sound",
         "I'm looking for a kids toothbrush for ages 6 to 12 that will help with teeth whitening and is easy to use",
+        "I need some cute heart-shaped glittery cupcake picks as a gift to bring to a baby shower",
     ],\
     title="WebShop",\
     article="<p style='padding-top:15px;text-align:center;'>To learn more about this project, check out the <a href='https://webshop-pnlp.github.io/' target='_blank'>project page</a>!</p>",\

predict.py

@@ -1,250 +0,0 @@
-import time, torch
-from transformers import BartTokenizer, BartForConditionalGeneration, AutoModel, AutoTokenizer
-
-from webshop_lite import dict_to_fake_html
-from predict_help import convert_dict_to_actions, convert_html_to_text, parse_results, parse_item_page, Page
-
-# Configurations
-DETAILED_OUTPUT = True
-BART_MODEL_PATH = 'webshop/il_search_bart'
-BERT_MODEL_PATH = 'webshop/il-rl-choice-bert-image_1'
-
-# load IL models
-bart_tokenizer = BartTokenizer.from_pretrained('facebook/bart-large')
-bart_model = BartForConditionalGeneration.from_pretrained(BART_MODEL_PATH)
-
-bert_tokenizer = AutoTokenizer.from_pretrained('bert-base-uncased', truncation_side='left')
-bert_tokenizer.add_tokens(['[button]', '[button_]', '[clicked button]', '[clicked button_]'], special_tokens=True)
-bert_model = AutoModel.from_pretrained(BERT_MODEL_PATH, trust_remote_code=True)
-
-def process_str(s):
-    s = s.lower().replace('"', '').replace("'", "").strip()
-    s = s.replace('[sep]', '[SEP]')
-    return s
-
-
-def process_goal(state):
-    state = state.lower().replace('"', '').replace("'", "")
-    state = state.replace('amazon shopping game\ninstruction:', '').replace('\n[button] search [button_]', '').strip()
-    if ', and price lower than' in state:
-        state = state.split(', and price lower than')[0]
-    return state
-
-
-def data_collator(batch):
-    state_input_ids, state_attention_mask, action_input_ids, action_attention_mask, sizes, labels, images = [], [], [], [], [], [], []
-    for sample in batch:
-        state_input_ids.append(sample['state_input_ids'])
-        state_attention_mask.append(sample['state_attention_mask'])
-        action_input_ids.extend(sample['action_input_ids'])
-        action_attention_mask.extend(sample['action_attention_mask'])
-        sizes.append(sample['sizes'])
-        labels.append(sample['labels'])
-        images.append(sample['images'])
-    max_state_len = max(sum(x) for x in state_attention_mask)
-    max_action_len = max(sum(x) for x in action_attention_mask)
-    return {
-        'state_input_ids': torch.tensor(state_input_ids)[:, :max_state_len],
-        'state_attention_mask': torch.tensor(state_attention_mask)[:, :max_state_len],
-        'action_input_ids': torch.tensor(action_input_ids)[:, :max_action_len],
-        'action_attention_mask': torch.tensor(action_attention_mask)[:, :max_action_len],
-        'sizes': torch.tensor(sizes),
-        'images': torch.tensor(images),
-        'labels': torch.tensor(labels),
-    }
-
-
-def bart_predict(input):
-    input_ids = bart_tokenizer(input)['input_ids']
-    input_ids = torch.tensor(input_ids).unsqueeze(0)
-    output = bart_model.generate(input_ids, max_length=512, num_return_sequences=5, num_beams=5)
-    return bart_tokenizer.batch_decode(output.tolist(), skip_special_tokens=True)[0]
-
-
-def bert_predict(obs, info, softmax=True):
-    valid_acts = info['valid']
-    assert valid_acts[0].startswith('click[')
-    state_encodings = bert_tokenizer(process_str(obs), max_length=512, truncation=True, padding='max_length')
-    action_encodings = bert_tokenizer(list(map(process_str, valid_acts)), max_length=512, truncation=True,  padding='max_length')
-    batch = {
-        'state_input_ids': state_encodings['input_ids'],
-        'state_attention_mask': state_encodings['attention_mask'],
-        'action_input_ids': action_encodings['input_ids'],
-        'action_attention_mask': action_encodings['attention_mask'],
-        'sizes': len(valid_acts),
-        'images': info['image_feat'].tolist(),
-        'labels': 0
-    }
-    batch = data_collator([batch])
-    outputs = bert_model(**batch)
-    if softmax:
-        idx = torch.multinomial(torch.nn.functional.softmax(outputs.logits[0], dim=0), 1)[0].item()
-    else:
-        idx = outputs.logits[0].argmax(0).item()
-    return valid_acts[idx]
-
-
-def predict(obs, info):
-    """
-    Given WebShop environment observation and info, predict an action.
-    """
-    valid_acts = info['valid']
-    if valid_acts[0].startswith('click['):
-        return bert_predict(obs, info)
-    else:
-        return "search[" + bart_predict(process_goal(obs)) + "]"
-
-def run_episode(goal, verbose=True):
-    """
-    Interact with amazon to find a product given input goal.
-    Input: text goal
-    Output: a url of found item on amazon.
-    """
-    obs = "Amazon Shopping Game\nInstruction:" + goal + "\n[button] search [button]"
-    info = {'valid': ['search[stuff]'], 'image_feat': torch.zeros(512)}
-    product_map = {}
-    title_to_asin_map = {}
-    search_results_cache = {}
-    visited_asins, clicked_options = set(), set()
-    sub_page_type, page_type, page_num = None, None, None
-    search_terms, prod_title, asin, num_prods, = None, None, None, None
-    options = {}
-    
-    for i in range(100):
-        # Run prediction
-        action = predict(obs, info)
-        if verbose:
-            print("====\n" + action)
-        
-        # Previous Page Type, Action -> Next Page Type
-        action_content = action[action.find("[")+1:action.find("]")]
-        prev_page_type = page_type
-        if action.startswith('search['):
-            page_type = Page.RESULTS
-            search_terms = action_content
-            page_num = 1
-        elif action.startswith('click['):
-            if action.startswith('click[item -'):
-                prod_title = action_content[len("item -"):].strip()
-                found = False
-                for key in title_to_asin_map:
-                    if prod_title == key:
-                        asin = title_to_asin_map[key]
-                        page_type = Page.ITEM_PAGE
-                        visited_asins.add(asin)
-                        found = True
-                        break
-                if not found:
-                    raise Exception("Product to click not found")
-                    
-            elif any(x.value in action for x in [Page.DESC, Page.FEATURES, Page.REVIEWS]):
-                page_type = Page.SUB_PAGE
-                sub_page_type = Page(action_content.lower())
-                
-            elif action == 'click[< prev]':
-                if sub_page_type is not None:
-                    page_type, sub_page_type = Page.ITEM_PAGE, None
-                elif prev_page_type == Page.ITEM_PAGE:
-                    page_type = Page.RESULTS
-                    options, clicked_options = {}, set()
-                elif prev_page_type == Page.RESULTS and page_num > 1:
-                    page_type = Page.RESULTS
-                    page_num -= 1
-                    
-            elif action == 'click[next >]':
-                page_type = Page.RESULTS
-                page_num += 1
-                
-            elif action.lower() == 'click[back to search]':
-                page_type = Page.SEARCH
-                
-            elif action == 'click[buy now]':
-                if DETAILED_OUTPUT:
-                    asin_url = f"https://www.amazon.com/dp/{asin}"
-                    return_value = "Product URL: " + asin_url
-                    if len(clicked_options) > 0:    
-                        options_str = ', '.join(list(clicked_options))
-                        return_value += "\nSelected Options: " + options_str
-                    return return_value
-                else:
-                    return asin
-            
-            elif prev_page_type == Page.ITEM_PAGE:
-                found = False
-                for opt_name, opt_values in product_map[asin]["options"].items():
-                    if action_content in opt_values:
-                        options[opt_name] = action_content
-                        page_type = Page.ITEM_PAGE
-                        clicked_options.add(action_content)
-                        found = True
-                        break
-                if not found:
-                    raise Exception("Unrecognized action: " + action)
-        else:
-            raise Exception("Unrecognized action:" + action)
-        
-        if verbose:
-            print(f"Parsing {page_type.value} page...")
-        
-        # URL -> Real HTML -> Dict of Info
-        if page_type == Page.RESULTS:
-            if search_terms in search_results_cache:
-                data = search_results_cache[search_terms]
-            else:
-                begin = time.time()
-                data = parse_results(search_terms, page_num, verbose)
-                end = time.time()
-                if verbose:
-                    print("Parsing search results took", end-begin, "seconds")
-
-                search_results_cache[search_terms] = data
-                num_prods = len(data)
-                for d in data:
-                    title_to_asin_map[d['Title']] = d['asin']
-        elif page_type == Page.ITEM_PAGE or page_type == Page.SUB_PAGE:
-            if asin in product_map:
-                if verbose:
-                    print("Loading cached item page for", asin)
-                data = product_map[asin]
-            else:
-                begin = time.time()
-                data = parse_item_page(asin, verbose)
-                end = time.time()
-                if verbose:
-                    print("Parsing item page took", end-begin, "seconds")
-                product_map[asin] = data
-        elif page_type == Page.SEARCH:
-            if verbose:
-                print("Executing search")
-            obs = "Amazon Shopping Game\nInstruction:" + goal + "\n[button] search [button]"
-            info = {'valid': ['search[stuff]'], 'image_feat': torch.zeros(512)}
-            continue
-        else:
-            raise Exception("Page of type `", page_type, "` not found")
-
-        # Dict of Info -> Fake HTML -> Text Observation
-        begin = time.time()
-        html_str = dict_to_fake_html(data, page_type, asin, sub_page_type, options, product_map, goal)
-        obs = convert_html_to_text(html_str, simple=False, clicked_options=clicked_options, visited_asins=visited_asins)
-        end = time.time()
-        if verbose:
-            print("[Page Info -> WebShop HTML -> Observation] took", end-begin, "seconds")
-
-        # Dict of Info -> Valid Action State (Info)
-        begin = time.time()
-        prod_arg = product_map if page_type == Page.ITEM_PAGE else data
-        info = convert_dict_to_actions(page_type, prod_arg, asin, page_num, num_prods)
-        end = time.time()
-        if verbose:
-            print("Extracting available actions took", end-begin, "seconds")
-        
-        if i == 99:
-            if DETAILED_OUTPUT:
-                asin_url = f"https://www.amazon.com/dp/{asin}"
-                return_value = "Product URL: " + asin_url
-                if len(clicked_options) > 0:    
-                    options_str = ', '.join(list(clicked_options))
-                    return_value += "\nSelected Options: " + options_str
-                return return_value
-            else:
-                return asin