Update app.py

app.py

@@ -1,59 +1,45 @@
 import streamlit as st
 import requests
+import asyncio
 import time
 from ast import literal_eval
+import urllib.parse
+from dacite import from_dict
+from together_web3.computer import LanguageModelInferenceRequest
+from together_web3.together import TogetherWeb3
+
+st.title("GPT-JT")
+if 'together_web3' not in st.session_state:
+    st.session_state.together_web3 = TogetherWeb3()
+if 'loop' not in st.session_state:
+    st.session_state.loop = asyncio.new_event_loop()
+async def _inference(prompt, max_tokens):
+    result = await st.session_state.together_web3.language_model_inference(
+        from_dict(
+            data_class=LanguageModelInferenceRequest,
+            data={
+                "model": "Together-gpt-JT-6B-v1",
+                "max_tokens": max_tokens,
+                "prompt": prompt,
+            }
+        ),
+    )
+    return result
 
 @st.cache
 def infer(prompt, 
           model_name, 
           max_new_tokens=10, 
-          temperature=0.0, 
+          temperature=0.0,
           top_p=1.0,
           num_completions=1,
           seed=42,
           stop="\n"):
-
-    model_name_map = {
-        "GPT-JT-6B-v1": "Together-gpt-JT-6B-v1",
-    }
-
-    my_post_dict = {
-        "type": "general",
-        "payload": {
-            "max_tokens": int(max_new_tokens),
-            "n": int(num_completions),
-            "temperature": float(temperature),
-            "top_p": float(top_p),
-            "model": model_name_map[model_name],
-            "prompt": [prompt],
-            "request_type": "language-model-inference",
-            "stop": stop.split(";"),
-            "best_of": 1,
-            "echo": False,
-            "seed": int(seed),
-            "prompt_embedding": False,
-        },
-        "returned_payload": {},
-        "status": "submitted",
-        "source": "dalle",
-    }
-    
-    job_id = requests.post("https://planetd.shift.ml/jobs", json=my_post_dict).json()['id']
-    
-    for i in range(100):
+    print("prompt", prompt)
+    response = st.session_state.loop.run_until_complete(_inference(prompt, int(max_new_tokens)))
+    print(response)
+    return response.choices[0].text
     
-        time.sleep(1)
-        
-        ret = requests.get(f"https://planetd.shift.ml/job/{job_id}", json={'id': job_id}).json()
-        
-        if ret['status'] == 'finished':
-            break
-        
-    return ret['returned_payload']['result']['inference_result'][0]['choices'][0]['text']
-    
-    
-st.title("GPT-JT")
-
 col1, col2 = st.columns([1, 3])
 
 with col1: