Update app.py

app.py

@@ -7,15 +7,22 @@ openchat_preprompt = (
 )
 
 def get_client(model: str):
-    InferenceAPIClient(model, token=os.getenv("HF_TOKEN", None))
+    if model == "togethercomputer/GPT-NeoXT-Chat-Base-20B":
+        return Client(os.getenv("OPENCHAT_API_URL"))
+    return InferenceAPIClient(model, token=os.getenv("HF_TOKEN", None))
+
 
 def get_usernames(model: str):
     """
     Returns:
         (str, str, str, str): pre-prompt, username, bot name, separator
     """
-    if model in ("OpenAssistant/oasst-sft-1-pythia-12b"):
+    if model in ("OpenAssistant/oasst-sft-1-pythia-12b", "OpenAssistant/oasst-sft-4-pythia-12b-epoch-3.5"):
         return "", "<|prompter|>", "<|assistant|>", "<|endoftext|>"
+    if model == "togethercomputer/GPT-NeoXT-Chat-Base-20B":
+        return openchat_preprompt, "<human>: ", "<bot>: ", "\n"
+    return "", "User: ", "Assistant: ", "\n"
+
 
 def predict(
     model: str,
@@ -52,7 +59,7 @@ def predict(
 
     partial_words = ""
 
-    if model in ("OpenAssistant/oasst-sft-1-pythia-12b"):
+    if model in ("OpenAssistant/oasst-sft-1-pythia-12b", "OpenAssistant/oasst-sft-4-pythia-12b-epoch-3.5"):
         iterator = client.generate_stream(
             total_inputs,
             typical_p=typical_p,
@@ -109,7 +116,7 @@ def radio_on_change(
     repetition_penalty,
     watermark,
 ):
-    if value in ("OpenAssistant/oasst-sft-1-pythia-12b"):
+    if value in ("OpenAssistant/oasst-sft-1-pythia-12b", "OpenAssistant/oasst-sft-4-pythia-12b-epoch-3.5"):
         typical_p = typical_p.update(value=0.2, visible=True)
         top_p = top_p.update(visible=False)
         top_k = top_k.update(visible=False)
@@ -117,6 +124,14 @@ def radio_on_change(
         disclaimer = disclaimer.update(visible=False)
         repetition_penalty = repetition_penalty.update(visible=False)
         watermark = watermark.update(False)
+    elif value == "togethercomputer/GPT-NeoXT-Chat-Base-20B":
+        typical_p = typical_p.update(visible=False)
+        top_p = top_p.update(value=0.25, visible=True)
+        top_k = top_k.update(value=50, visible=True)
+        temperature = temperature.update(value=0.6, visible=True)
+        repetition_penalty = repetition_penalty.update(value=1.01, visible=True)
+        watermark = watermark.update(False)
+        disclaimer = disclaimer.update(visible=True)
     else:
         typical_p = typical_p.update(visible=False)
         top_p = top_p.update(value=0.95, visible=True)
@@ -136,18 +151,14 @@ def radio_on_change(
     )
 
 
-title = """<h1 align="center">LLM Chat</h1>"""
-description = """LLM Chat predložak:
-```
-Mjenjajući predložak.
+title = """<h1 align="center">xChat</h1>"""
+description = """
 """
 
 text_generation_inference = """
-<div align="center">Pokrenuto od: <a href=https://xnet.ba/>xnet.ba</a></div>
 """
 
 openchat_disclaimer = """
-<div align="center">Checkout the official <a href=https://xnet.ba>xChat app</a> for the full experience.</div>
 """
 
 with gr.Blocks(
@@ -162,6 +173,7 @@ with gr.Blocks(
             choices=[
                 "OpenAssistant/oasst-sft-4-pythia-12b-epoch-3.5",
                 "OpenAssistant/oasst-sft-1-pythia-12b",
+                "togethercomputer/GPT-NeoXT-Chat-Base-20B",
             ],
             label="Model",
             interactive=True,