Update app.py

app.py

@@ -1,3 +1,4 @@
+
 import streamlit as st
 import requests
 import time
@@ -18,11 +19,9 @@ def infer(
     seed=42,
     stop="\n"
 ):
-    
     model_name_map = {
         "GPT-JT-6B-v1": "Together-gpt-JT-6B-v1",
     }
-    
     max_new_tokens = int(max_new_tokens)
     num_completions = int(num_completions)
     temperature = float(temperature)
@@ -39,21 +38,41 @@ def infer(
         temperature = 0.01
 
     my_post_dict = {
-        "model": "Together-gpt-JT-6B-v1",
-        "prompt": prompt,
-        "top_p": top_p,
-        "top_k": top_k,
-        "temperature": temperature,
-        "max_tokens": max_new_tokens,
-        "stop": stop,
+        "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,
+            "best_of": 1,
+            "echo": False,
+            "seed": int(seed),
+            "prompt_embedding": False,
+        },
+        "returned_payload": {},
+        "status": "submitted",
+        "source": "dalle",
     }
-    response = requests.get("https://staging.together.xyz/api/inference", params=my_post_dict).json()
-    generated_text = response['output']['choices'][0]['text']
+    
+    job_id = requests.post("https://planetd.shift.ml/jobs", json=my_post_dict).json()['id']
+    
+    for i in range(100):
+    
+        time.sleep(0.5)
+        ret = requests.get(f"https://planetd.shift.ml/job/{job_id}", json={'id': job_id}).json()
+        if ret['status'] == 'finished':
+            break
+
+    generated_text = ret['returned_payload']['result']['inference_result'][0]['choices'][0]['text']
     
     for stop_word in stop:
         if stop_word in generated_text:
             generated_text = generated_text[:generated_text.find(stop_word)]
-            
+
     st.session_state.updated = True
     
     return generated_text
@@ -84,7 +103,7 @@ Label:'''
 
 Input: bear, honey
 Story: Once upon a time,'''
-        st.session_state.temperature = "0.8"
+        st.session_state.temperature = "0.0"
         st.session_state.top_p = "0.9"
         st.session_state.max_new_tokens = "100"
     
@@ -101,7 +120,7 @@ def main():
         st.session_state.prompt = "Please answer the following question:\n\nQuestion: In which country is Zurich located?\nAnswer:"
 
     if 'temperature' not in st.session_state:
-        st.session_state.temperature = "0.8"
+        st.session_state.temperature = "0.0"
 
     if 'top_p' not in st.session_state:
         st.session_state.top_p = "1.0"
@@ -164,4 +183,4 @@ def main():
             generated_area.markdown("<b>" + to_md(prompt) + "</b>" + to_md(report_text), unsafe_allow_html=True)
         
 if __name__ == '__main__':
-    main()
+    main()