Update app.py

app.py

@@ -1,8 +1,10 @@
 
+
 import streamlit as st
 import requests
 import time
 from ast import literal_eval
+from datetime import datetime
 
 def to_md(text):
     return text.replace("\n", "<br />")
@@ -38,41 +40,23 @@ def infer(
         temperature = 0.01
 
     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,
-            "best_of": 1,
-            "echo": False,
-            "seed": int(seed),
-            "prompt_embedding": False,
-        },
-        "returned_payload": {},
-        "status": "submitted",
-        "source": "dalle",
+        "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,
     }
-    
-    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']
+    print(f"send: {datetime.now()}")
+    response = requests.get("https://staging.together.xyz/api/inference", params=my_post_dict).json()
+    generated_text = response['output']['choices'][0]['text']
+    print(f"recv: {datetime.now()}")
     
     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
@@ -120,7 +104,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.1"
+        st.session_state.temperature = "0.8"
 
     if 'top_p' not in st.session_state:
         st.session_state.top_p = "1.0"