update

app.py

@@ -3,86 +3,162 @@ import requests
 import time
 from ast import literal_eval
 
-@st.cache
-def infer(prompt, 
-          model_name, 
-          max_new_tokens=10, 
-          temperature=0.0, 
-          top_p=1.0,
-          num_completions=1,
-          seed=42,
-          stop="\n"):
+    
+def infer(
+    prompt, 
+    model_name, 
+    max_new_tokens=10, 
+    temperature=0.0, 
+    top_p=1.0,
+    top_k=40,
+    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']
+    max_new_tokens = int(max_new_tokens)
+    num_completions = int(num_completions)
+    temperature = float(temperature)
+    top_p = float(top_p)
+    stop = stop.split(";")
+    seed = seed
     
-    for i in range(100):
+    assert 0 <= max_new_tokens <= 256
+    assert 1 <= num_completions <= 5
+    assert 0.0 <= temperature <= 10.0
+    assert 0.0 <= top_p <= 1.0
     
-        time.sleep(1)
+    if temperature == 0.0:
+        temperature = 1.0
+        top_k = 1
+    
+    result = await st.session_state.together_web3.language_model_inference(
+        from_dict(
+            data_class=LanguageModelInferenceRequest,
+            data={
+                "model": model_name_map[model_name],
+                "max_tokens": max_new_tokens,
+                "prompt": prompt,
+                "n": num_completions,
+                "temperature": temperature,
+                "top_k": top_k,
+                "top_p": top_p,
+                "stop": stop,
+                "seed": seed,
+                "echo": False,
+            }
+        ),
+    )
+    
+    generated_text = result.choices[0].text
+    
+    for stop_word in stop:
+        if stop_word in result:
+            generated_text = generated_text[:generated_text.find(stop_word)]
+    
+    return generated_text
+
+def set_preset():
+    if st.session_state.preset == "Classification":
         
-        ret = requests.get(f"https://planetd.shift.ml/job/{job_id}", json={'id': job_id}).json()
+        st.session_state.prompt = '''Please classify the given sentence.
+Possible labels:
+1. <label_0>
+2. <label_1>
+
+Input: <sentence_0>
+Label: <label_0>
+
+Input: <sentence_1>
+Label:'''
+        st.session_state.temperature = "0.0"
+        st.session_state.top_p = "1.0"
         
-        if ret['status'] == 'finished':
-            break
+    elif st.session_state.preset == "Generation":
         
-    return ret['returned_payload']['result']['inference_result'][0]['choices'][0]['text']
+        st.session_state.prompt = '''Please write a story given keywords.
+
+Input: bear, honey
+Story:'''
+        st.session_state.temperature = "1.0"
+        st.session_state.top_p = "0.5"
     
+    else:
+        pass
     
-st.title("GPT-JT")
-
-col1, col2 = st.columns([1, 3])
-
-with col1:
-    model_name = st.selectbox("Model", ["GPT-JT-6B-v1"])
-    max_new_tokens = st.text_input('Max new tokens', "10")
-    temperature = st.text_input('temperature', "0.0")
-    top_p = st.text_input('top_p', "1.0")
-    num_completions = st.text_input('num_completions (only the best one will be returend)', "1")
-    stop = st.text_input('stop, split by;', r'\n')
-    seed = st.text_input('seed', "42")
-
-with col2:
-    s_example = "Please answer the following question:\n\nQuestion: Where is Zurich?\nAnswer:"
-    prompt = st.text_area(
-        "Prompt",
-        value=s_example,
-        max_chars=4096,
-        height=400,
-    )
-        
-    generated_area = st.empty()
-    generated_area.text("(Generate here)")
     
-    button_submit = st.button("Submit")
+def main():
+
+    if 'preset' not in st.session_state:
+        st.session_state.preset = "Classification"
+
+    if 'prompt' not in st.session_state:
+        st.session_state.prompt = "Please answer the following question:\n\nQuestion: Where is Zurich?\nAnswer:"
+
+    if 'temperature' not in st.session_state:
+        st.session_state.temperature = "0.0"
 
-    if button_submit:
-        generated_area.text(prompt)
-        report_text = infer(
-            prompt, model_name=model_name, max_new_tokens=max_new_tokens, temperature=temperature, top_p=top_p,
-            num_completions=num_completions, seed=seed, stop=literal_eval("'''"+stop+"'''"),
+    if 'top_p' not in st.session_state:
+        st.session_state.top_p = "1.0"
+
+    if 'top_k' not in st.session_state:
+        st.session_state.top_k = "40"
+        
+    if 'together_web3' not in st.session_state:
+        st.session_state.together_web3 = TogetherWeb3()
+        
+
+    st.title("GPT-JT")
+
+    col1, col2 = st.columns([1, 3])
+
+    with col1:
+        model_name = st.selectbox("Model", ["GPT-JT-6B-v1"])
+        max_new_tokens = st.text_input('Max new tokens', "10")
+        temperature = st.text_input('temperature', st.session_state.temperature)
+        top_k = st.text_input('top_k', st.session_state.top_k)
+        top_p = st.text_input('top_p', st.session_state.top_p)
+        # num_completions = st.text_input('num_completions (only the best one will be returend)', "1")
+        num_completions = "1"
+        stop = st.text_input('stop, split by;', r'\n')
+        # seed = st.text_input('seed', "42")
+        seed = "42"
+
+    with col2:
+
+        preset = st.radio(
+            "Recommended Configurations", 
+            ('Classification', 'Generation'), 
+            on_change=set_preset,
+            key="preset",
+            horizontal=True
         )
-        generated_area.text(prompt + report_text)
+
+        prompt = st.text_area(
+            "Prompt",
+            value=st.session_state.prompt,
+            max_chars=4096,
+            height=400,
+        )
+
+        generated_area = st.empty()
+        generated_area.text("(Generate here)")
+
+        button_submit = st.button("Submit")
+
+        if button_submit:
+            generated_area.text(prompt)
+            report_text = infer(
+                prompt, model_name=model_name, max_new_tokens=max_new_tokens, temperature=temperature, top_p=top_p, top_k=top_k,
+                num_completions=num_completions, seed=seed, stop=literal_eval("'''"+stop+"'''"),
+            )
+            generated_area.text(prompt + report_text)
+
+        
+        
+if __name__ == '__main__':
+    main()