Add txt2txt clients

0_🏠_Home.py

@@ -39,10 +39,9 @@ st.html("""
     <h1 style="padding: 0; margin-bottom: 0.5rem">API Inference</h1>
     <span class="pro-badge">PRO</span>
 </div>
-<p>Run inference on API endpoints. Hugging Face for now; more coming soon!</p>
+<p>Inference on Huggingface, Perplexity, and Fal ⚡</p>
 """)
 
-# TODO: categorize tasks by service (SAI, FAL, etc)
 # content
 st.markdown("## Tasks")
 st.page_link("pages/1_💬_Text_Generation.py", label="Text Generation", icon="💬")

README.md

@@ -46,6 +46,14 @@ STREAMLIT_SERVER_RUN_ON_SAVE=false
 STREAMLIT_BROWSER_SERVER_ADDRESS=adamelliotfields-api-inference.hf.space
 ```
 
+## Secrets
+
+```bash
+FAL_KEY=...
+HF_TOKEN=...
+PPLX_API_KEY=...
+```
+
 ## Installation
 
 ```sh

lib/__init__.py

@@ -1,5 +1,12 @@
-from .api import HuggingFaceTxt2ImgAPI
+from .api import HuggingFaceTxt2ImgAPI, HuggingFaceTxt2TxtAPI, PerplexityTxt2TxtAPI
 from .config import Config
-from .presets import Presets
+from .presets import ModelPresets, ServicePresets
 
-__all__ = ["Config", "HuggingFaceTxt2ImgAPI", "Presets"]
+__all__ = [
+    "Config",
+    "HuggingFaceTxt2ImgAPI",
+    "HuggingFaceTxt2TxtAPI",
+    "ModelPresets",
+    "PerplexityTxt2TxtAPI",
+    "ServicePresets",
+]

lib/api.py

@@ -2,15 +2,63 @@ import io
 from abc import ABC, abstractmethod
 
 import requests
+import streamlit as st
+from openai import APIError, OpenAI
 from PIL import Image
 
 
+class Txt2TxtAPI(ABC):
+    @abstractmethod
+    def generate_text(self, model, parameters, **kwargs):
+        pass
+
+
 class Txt2ImgAPI(ABC):
     @abstractmethod
     def generate_image(self, model, prompt, parameters, **kwargs):
         pass
 
 
+class HuggingFaceTxt2TxtAPI(Txt2TxtAPI):
+    def __init__(self, api_key):
+        self.api_key = api_key
+
+    def generate_text(self, model, parameters, **kwargs):
+        if not self.api_key:
+            return "API Key is required."
+        client = OpenAI(
+            api_key=self.api_key,
+            base_url=f"https://api-inference.huggingface.co/models/{model}/v1",
+        )
+        try:
+            stream = client.chat.completions.create(stream=True, model=model, **parameters, **kwargs)
+            return st.write_stream(stream)
+        except APIError as e:
+            return e.message
+        except Exception as e:
+            return str(e)
+
+
+class PerplexityTxt2TxtAPI(Txt2TxtAPI):
+    def __init__(self, api_key):
+        self.api_key = api_key
+
+    def generate_text(self, model, parameters, **kwargs):
+        if not self.api_key:
+            return "API Key is required."
+        client = OpenAI(
+            api_key=self.api_key,
+            base_url="https://api.perplexity.ai",
+        )
+        try:
+            stream = client.chat.completions.create(stream=True, model=model, **parameters, **kwargs)
+            return st.write_stream(stream)
+        except APIError as e:
+            return e.message
+        except Exception as e:
+            return str(e)
+
+
 # essentially the same as huggingface_hub's inference client
 class HuggingFaceTxt2ImgAPI(Txt2ImgAPI):
     def __init__(self, token):

lib/config.py

@@ -19,12 +19,24 @@ Config = SimpleNamespace(
         "7:9": (896, 1152),
         "4:7": (768, 1344),
     },
-    TXT2TXT_DEFAULT_MODEL=4,
-    TXT2TXT_MODELS=[
-        "codellama/codellama-34b-instruct-hf",
-        "meta-llama/llama-2-13b-chat-hf",
-        "meta-llama/meta-llama-3.1-405b-instruct-fp8",
-        "mistralai/mistral-7b-instruct-v0.2",
-        "nousresearch/nous-hermes-2-mixtral-8x7b-dpo",
-    ],
+    TXT2TXT_DEFAULT_SYSTEM="You are a helpful assistant. Be precise and concise.",
+    TXT2TXT_DEFAULT_MODEL={
+        "Huggingface": 4,
+        "Perplexity": 3,
+    },
+    TXT2TXT_MODELS={
+        "Huggingface": [
+            "codellama/codellama-34b-instruct-hf",
+            "meta-llama/llama-2-13b-chat-hf",
+            "meta-llama/meta-llama-3.1-405b-instruct-fp8",
+            "mistralai/mistral-7b-instruct-v0.2",
+            "nousresearch/nous-hermes-2-mixtral-8x7b-dpo",
+        ],
+        "Perplexity": [
+            "llama-3.1-sonar-small-128k-chat",
+            "llama-3.1-sonar-large-128k-chat",
+            "llama-3.1-sonar-small-128k-online",
+            "llama-3.1-sonar-large-128k-online",
+        ],
+    },
 )

lib/presets.py

@@ -1,6 +1,22 @@
 from types import SimpleNamespace
 
-Presets = SimpleNamespace(
+ServicePresets = SimpleNamespace(
+    Huggingface={
+        # every service has model and system messages
+        "frequency_penalty": 0.0,
+        "frequency_penalty_min": -2.0,
+        "frequency_penalty_max": 2.0,
+        "parameters": ["max_tokens", "temperature", "frequency_penalty", "seed"],
+    },
+    Perplexity={
+        "frequency_penalty": 1.0,
+        "frequency_penalty_min": 1.0,
+        "frequency_penalty_max": 2.0,
+        "parameters": ["max_tokens", "temperature", "frequency_penalty"],
+    },
+)
+
+ModelPresets = SimpleNamespace(
     FLUX_1_DEV={
         "name": "FLUX.1 Dev",
         "num_inference_steps": 30,

pages/1_💬_Text_Generation.py

@@ -2,23 +2,20 @@ import os
 from datetime import datetime
 
 import streamlit as st
-from openai import APIError, OpenAI
 
-from lib import Config
+from lib import Config, HuggingFaceTxt2TxtAPI, PerplexityTxt2TxtAPI, ServicePresets
 
-# TODO: key input and store in cache_data
-# api key
-HF_TOKEN = os.environ.get("HF_TOKEN")
+HF_TOKEN = os.environ.get("HF_TOKEN") or None
+PPLX_API_KEY = os.environ.get("PPLX_API_KEY") or None
 
 
-# TODO: eventually support different APIs like OpenAI and Perplexity
 @st.cache_resource
-def get_chat_client(api_key, model):
-    client = OpenAI(
-        api_key=api_key,
-        base_url=f"https://api-inference.huggingface.co/models/{model}/v1",
-    )
-    return client
+def get_txt2txt_api(service="Huggingface", api_key=None):
+    if service == "Huggingface":
+        return HuggingFaceTxt2TxtAPI(api_key)
+    if service == "Perplexity":
+        return PerplexityTxt2TxtAPI(api_key)
+    return None
 
 
 # config
@@ -29,6 +26,9 @@ st.set_page_config(
 )
 
 # initialize state
+if "txt2txt_running" not in st.session_state:
+    st.session_state.txt2txt_running = False
+
 if "txt2txt_messages" not in st.session_state:
     st.session_state.txt2txt_messages = []
 
@@ -38,52 +38,83 @@ if "txt2txt_prompt" not in st.session_state:
 # sidebar
 st.logo("logo.svg")
 st.sidebar.header("Settings")
+service = st.sidebar.selectbox(
+    "Service",
+    options=["Huggingface", "Perplexity"],
+    index=0,
+    disabled=st.session_state.txt2txt_running,
+)
+
+# hide key input if environment variables are set
+if (service == "Huggingface" and HF_TOKEN is None) or (service == "Perplexity" and PPLX_API_KEY is None):
+    api_key = st.sidebar.text_input(
+        "API Key",
+        value="",
+        type="password",
+        help="Cleared on page refresh",
+        disabled=st.session_state.txt2txt_running,
+    )
+
 model = st.sidebar.selectbox(
     "Model",
-    placeholder="Select a model",
-    format_func=lambda x: x.split("/")[1],
-    index=Config.TXT2TXT_DEFAULT_MODEL,
-    options=Config.TXT2TXT_MODELS,
-)
-max_tokens = st.sidebar.slider(
-    "Max Tokens",
-    min_value=512,
-    max_value=4096,
-    value=512,
-    step=128,
-    help="Maximum number of tokens to generate (default: 512)",
-)
-temperature = st.sidebar.slider(
-    "Temperature",
-    min_value=0.0,
-    max_value=2.0,
-    value=1.0,
-    step=0.1,
-    help="Used to modulate the next token probabilities (default: 1.0)",
-)
-frequency_penalty = st.sidebar.slider(
-    "Frequency Penalty",
-    min_value=-2.0,
-    max_value=2.0,
-    value=0.0,
-    step=0.1,
-    help="Penalize new tokens based on their existing frequency in the text (default: 0.0)",
-)
-seed = st.sidebar.number_input(
-    "Seed",
-    min_value=-1,
-    max_value=(1 << 53) - 1,
-    value=-1,
-    help="Make a best effort to sample deterministically (default: -1)",
+    format_func=lambda x: x.split("/")[1] if service == "Huggingface" else x,
+    index=Config.TXT2TXT_DEFAULT_MODEL[service],
+    options=Config.TXT2TXT_MODELS[service],
+    disabled=st.session_state.txt2txt_running,
 )
 system = st.sidebar.text_area(
     "System Message",
-    value="You are a helpful assistant. Be precise and concise.",
+    value=Config.TXT2TXT_DEFAULT_SYSTEM,
+    disabled=st.session_state.txt2txt_running,
 )
 
+# build parameters from preset
+parameters = {}
+preset = getattr(ServicePresets, service)
+for param in preset["parameters"]:
+    if param == "max_tokens":
+        parameters[param] = st.sidebar.slider(
+            "Max Tokens",
+            min_value=512,
+            max_value=4096,
+            value=512,
+            step=128,
+            help="Maximum number of tokens to generate (default: 512)",
+            disabled=st.session_state.txt2txt_running,
+        )
+    if param == "temperature":
+        parameters[param] = st.sidebar.slider(
+            "Temperature",
+            min_value=0.0,
+            max_value=2.0,
+            value=1.0,
+            step=0.1,
+            help="Used to modulate the next token probabilities (default: 1.0)",
+            disabled=st.session_state.txt2txt_running,
+        )
+    if param == "frequency_penalty":
+        parameters[param] = st.sidebar.slider(
+            "Frequency Penalty",
+            min_value=preset["frequency_penalty_min"],
+            max_value=preset["frequency_penalty_max"],
+            value=preset["frequency_penalty"],
+            step=0.1,
+            help="Penalize new tokens based on their existing frequency in the text (default: 0.0)",
+            disabled=st.session_state.txt2txt_running,
+        )
+    if param == "seed":
+        parameters[param] = st.sidebar.number_input(
+            "Seed",
+            min_value=-1,
+            max_value=(1 << 53) - 1,
+            value=-1,
+            help="Make a best effort to sample deterministically (default: -1)",
+            disabled=st.session_state.txt2txt_running,
+        )
+
 # random seed
-if seed < 0:
-    seed = int(datetime.now().timestamp() * 1e6) % (1 << 53)
+if parameters.get("seed", 0) < 0:
+    parameters["seed"] = int(datetime.now().timestamp() * 1e6) % (1 << 53)
 
 # heading
 st.html("""
@@ -136,7 +167,10 @@ else:
     button_container = None
 
 # chat input
-if prompt := st.chat_input("What would you like to know?"):
+if prompt := st.chat_input(
+    "What would you like to know?",
+    on_submit=lambda: setattr(st.session_state, "txt2txt_running", True),
+):
     st.session_state.txt2txt_prompt = prompt
 
 if st.session_state.txt2txt_prompt:
@@ -149,24 +183,19 @@ if st.session_state.txt2txt_prompt:
     messages = [{"role": "system", "content": system}]
     messages.extend([{"role": m["role"], "content": m["content"]} for m in st.session_state.txt2txt_messages])
     messages.append({"role": "user", "content": st.session_state.txt2txt_prompt})
+    parameters["messages"] = messages
 
     with st.chat_message("assistant"):
-        try:
-            client = get_chat_client(HF_TOKEN, model)
-            stream = client.chat.completions.create(
-                frequency_penalty=frequency_penalty,
-                temperature=temperature,
-                max_tokens=max_tokens,
-                messages=messages,
-                model=model,
-                stream=True,
-                seed=seed,
-            )
-            response = st.write_stream(stream)
-        except APIError as e:
-            response = e.message
-        except Exception as e:
-            response = str(e)
+        # allow environment variables in development for convenience
+        if service == "Huggingface" and HF_TOKEN is not None:
+            key = HF_TOKEN
+        elif service == "Perplexity" and PPLX_API_KEY is not None:
+            key = PPLX_API_KEY
+        else:
+            key = api_key
+        api = get_txt2txt_api(service, key)
+        response = api.generate_text(model, parameters)
+        st.session_state.txt2txt_running = False
 
     st.session_state.txt2txt_messages.append({"role": "user", "content": st.session_state.txt2txt_prompt})
     st.session_state.txt2txt_messages.append({"role": "assistant", "content": response})

pages/2_🎨_Text_to_Image.py

@@ -3,7 +3,7 @@ from datetime import datetime
 
 import streamlit as st
 
-from lib import Config, HuggingFaceTxt2ImgAPI, Presets
+from lib import Config, HuggingFaceTxt2ImgAPI, ModelPresets
 
 # TODO: key input and store in cache_data
 # TODO: API dropdown; changes available models
@@ -11,9 +11,9 @@ HF_TOKEN = os.environ.get("HF_TOKEN")
 API_URL = "https://api-inference.huggingface.co/models"
 HEADERS = {"Authorization": f"Bearer {HF_TOKEN}", "X-Wait-For-Model": "true", "X-Use-Cache": "false"}
 PRESET_MODEL = {
-    "black-forest-labs/flux.1-dev": Presets.FLUX_1_DEV,
-    "black-forest-labs/flux.1-schnell": Presets.FLUX_1_SCHNELL,
-    "stabilityai/stable-diffusion-xl-base-1.0": Presets.STABLE_DIFFUSION_XL,
+    "black-forest-labs/flux.1-dev": ModelPresets.FLUX_1_DEV,
+    "black-forest-labs/flux.1-schnell": ModelPresets.FLUX_1_SCHNELL,
+    "stabilityai/stable-diffusion-xl-base-1.0": ModelPresets.STABLE_DIFFUSION_XL,
 }