Add `env` input

app.py

@@ -13,13 +13,17 @@ from predict_help import (
 
 ENVIRONMENTS = ['amazon', 'webshop', 'ebay']
 
+# IL+RL: 'webshop/il-rl-choice-bert-image_1'
+# IL: 'webshop/il-choice-bert-image_0'
+BERT_MODEL_PATH = 'webshop/il-choice-bert-image_0'
+
 # 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)
+bert_model = AutoModel.from_pretrained(BERT_MODEL_PATH, trust_remote_code=True)
 
 def process_str(s):
     s = s.lower().replace('"', '').replace("'", "").strip()
@@ -98,12 +102,13 @@ def predict(obs, info):
     else:
         return "search[" + bart_predict(process_goal(obs)) + "]"
 
-def run_episode(goal, verbose=True, env='amazon'):
+def run_episode(goal, env, verbose=True):
     """
     Interact with amazon to find a product given input goal.
     Input: text goal
     Output: a url of found item on amazon.
     """
+    env = env.lower()
     if env not in ENVIRONMENTS:
         print(f"[ERROR] Environment {env} not recognized")
         
@@ -289,14 +294,17 @@ def run_episode(goal, verbose=True, env='amazon'):
             return return_value
 
 gr.Interface(fn=run_episode,\
-    inputs=gr.inputs.Textbox(lines=7, label="Input Text"),\
+    inputs=[
+        gr.inputs.Textbox(lines=7, label="Input Text"),
+        gr.inputs.Radio(['Amazon', 'eBay', 'WebShop'], type="value", default="Amazon", label='Environment')
+    ],\
     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",
+        ["Please select a 1 pound, certified organic sea salt shaker in the flavor triple blend flakes, and price lower than 40.00 dollars", "Amazon"],
+        ["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", "Amazon"],
+        ["I'm trying to find white bluetooth speakers that are not only water resistant but also come with stereo sound", "Amazon"],
+        ["I need some cute heart-shaped glittery cupcake picks as a gift to bring to a baby shower", "eBay"],
+        ["I'm looking for a kids toothbrush for ages 6 to 12 that will help with teeth whitening and is easy to use", "WebShop"],
     ],\
     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>",\