Merge pull request #3 from east-and-west-magic/feature-three-pipeline

app.py

@@ -1,5 +1,6 @@
 import requests
 import json
+import os
 from datetime import datetime
 from zoneinfo import ZoneInfo
 
@@ -8,16 +9,16 @@ from gradio_client import Client
 from extract import extract
 
 functionality_shared = [
-    'games',
-    'greet',
-    'tokenization',
+    "games",
+    "greet",
+    "tokenization",
     "upload game",
     "download game",
 ]
 
 game_list = [
-    'angry birds',
-    'snake',
+    "angry birds",
+    "snake",
 ]
 
 functionality_list = [
@@ -28,13 +29,28 @@ functionality_list = [
 ]
 
 t = datetime.now()
-t = t.astimezone(ZoneInfo('Asia/Shanghai'))
+t = t.astimezone(ZoneInfo("Asia/Shanghai"))
 print(f"[Beijing]: {t.replace(microsecond=0)}")
-t = t.astimezone(ZoneInfo('America/Los_Angeles'))
+t = t.astimezone(ZoneInfo("America/Los_Angeles"))
 print(f"[Seattle]: {t.replace(microsecond=0)}")
 
 # assert False
 
+
+spaces = {
+    "b_demo_hf": "stevez-proxy",
+    "pgdemo2": "stevez-proxy2",
+    "pgdemo3": "stevez-proxy3",
+}
+
+identity = os.environ.get("identity")
+if not identity:
+    identity = "b_demo_hf"
+space = "stevez-proxy"
+if identity in spaces:
+    space = spaces[identity]
+
+
 def run(hf_token, service, game, functionality, nlp_command):
     """
     event handler
@@ -43,22 +59,22 @@ def run(hf_token, service, game, functionality, nlp_command):
     # reuse hf_token field as json string
     token, user_name, redirect, _ = extract(hf_token)
     if user_name is None:
-        user_name = '__fake__'
+        user_name = "__fake__"
 
     # redirect all traffic to the proxy sever
-    space = 'stevez-proxy'
     if redirect is not None:
         space = redirect
-        
-    if token is None or token == '':
-        return 'please specify hf token'
+    url = f"https://{space}.hf.space"
+
+    if token is None or token == "":
+        return "please specify hf token"
 
     if service not in functionality_shared[1:]:
         if game is None:
             return "please specify which game"
         if functionality is None:
             return "please choose the AI functionality"
-        if functionality == 'AI':
+        if functionality == "AI":
             if nlp_command in ["", None]:
                 return "please make sure the command is not empty"
 
@@ -67,13 +83,18 @@ def run(hf_token, service, game, functionality, nlp_command):
             pass
 
         client = Client(
-            f'https://{space}.hf.space',
+            url,
             hf_token=token,
             verbose=False,
         )
         res = client.predict(
-            user_name, service, game, functionality, nlp_command, nlp_command,
-            api_name='/predict',
+            user_name,
+            service,
+            game,
+            functionality,
+            nlp_command,
+            nlp_command,
+            api_name="/predict",
         )
         # if service in [
         #     "upload game",
@@ -87,31 +108,35 @@ def run(hf_token, service, game, functionality, nlp_command):
         res2 = json.loads(res[0])
         return json.dumps(res2, indent=4)
     except Exception as e:
-        return f"{type(e)}, {str(e)}. \nyou may want to make sure your hf_token is correct"
-    
+        return (
+            f"{type(e)}, {str(e)}. \nyou may want to make sure your hf_token is correct"
+        )
+
+
 demo = gr.Interface(
-    fn=run, 
+    fn=run,
     inputs=[
         "text",
         gr.Radio(
             functionality_shared,
             value=functionality_shared[0],
-            info = "Shared services",
-            ),
+            info="Shared services",
+        ),
         gr.Radio(
             game_list,
-            value=game_list[1],        
-            info = "Which game you want the AI to support?",
-            ),
+            value=game_list[1],
+            info="Which game you want the AI to support?",
+        ),
         gr.Radio(
             functionality_list,
             value=functionality_list[0],
             # label = "What do you want to do?",
-            info = "What functionality?",
-            ),
-        "text"],
+            info="What functionality?",
+        ),
+        "text",
+    ],
     outputs="text",
-    title = "Demo",
+    title="Demo",
 )
 
 demo.launch()