HF api + Cohere api + TxtToTxt + POST w/ retry

api/base_api_provider.py

@@ -1,3 +1,4 @@
+import json
 from abc import abstractmethod, ABCMeta
 from pathlib import Path
 from typing import Any, Dict, Optional
@@ -11,7 +12,7 @@
 
 
 class BaseAPIProvider(metaclass=ABCMeta):
-    def __init__(self, storage_dir: str = "~/.chisel"):
+    def __init__(self, storage_dir: str = "~/chisel"):
         self.storage = LocalFS(storage_dir)
         s = requests.Session()
         retries = Retry(
@@ -57,6 +58,15 @@ def _download_img_from_url(
             full_path = self.storage.stream_to_tmp(r, filename, ext)
         return full_path
 
+    def _post_with_retry(
+        self,
+        url: str,
+        headers: Dict[str, str],
+        payload: Any,
+    ) -> Dict[str, Any]:
+        data = json.dumps(payload)
+        return self.requests_session.post(url, headers=headers, data=data)
+
 
 class APIResult(object):
     def __init__(self) -> None:

api/hf.py

@@ -1,9 +1,8 @@
 import json
-from pathlib import Path
-from typing import Any, Dict, List, Optional
+import os
+from typing import Any, Dict, List
 
-import requests
-import openai
+from PIL import Image
 
 from chisel.api.base_api_provider import (
     APIResult, BaseAPIProvider
@@ -20,42 +19,30 @@ def __init__(self) -> None:
     def _process_results(self, response) -> List[Any]:
         api_result = APIResult()
 
-        results = response.get('data', None)
-        if results is None:
-            return api_result
-
-        for i, result in enumerate(results):
-            url = result.get("url", None)
-            if url is None:
-                continue
-            local_filename = self._download_img_from_url(url, ext=".png")
-            api_result.add(local_filename, remote_url=url)
-        return api_result
+        if response.headers['Content-Type'] == 'image/jpeg':
+            data = response.content
+            full_path = self.storage.write_to_tmp(data, ext=".png")
+            return Image.open(full_path)
+        else:
+            data = json.loads(response.content.decode("utf-8"))
+            self.storage.write_to_tmp(data, ext=".txt")
+            return data
 
 
 class HFInference(HF):
     def __init__(self) -> None:
         super().__init__()
-        self.params: Dict[str, Any] = {
-            "width": 512,
-            "samples": 1,
-        }
-
-    def _query(self, payload):
-        data = json.dumps(payload)
-        response = requests.request("POST", API_URL, headers=headers, data=data)
-        return json.loads(response.content.decode("utf-8"))
 
     def run(self, inp: Any, params: Dict[str, str] = None) -> Any:
-        if not isinstance(inp, str):
-            raise ValueError("Expected inp to be a str (prompt).")
-
         model_id = params.get("model_id", None)
-        API_URL = f"https://api-inference.huggingface.co/models/{model_id}"
+        txt_to_img = params.get("txt_to_img", False)
 
-        print(self.api_key)
+        api_url = f"https://api-inference.huggingface.co/models/{model_id}"
 
-        headers = {"Authorization": f"Bearer {self.api_key}"}
+        headers = {
+            "Authorization": f"Bearer {self.api_key}",
+            'Content-Type': 'application/json',
+        }
 
-        data = query("Can you please let us know more details about your ")
+        response = self._post_with_retry(api_url, headers, inp)
         return self._process_results(response)

ops/__init__.py

@@ -1,3 +1,4 @@
+from chisel.ops.txt_to_txt import TxtToTxt
 from chisel.ops.txt_to_img import TxtToImg
 from chisel.ops.img_to_img import ImgToImg
 from chisel.ops.img_edit import ImgEdit

ops/provider.py

@@ -7,3 +7,4 @@ class Provider(str, Enum):
     OPENAI = "openai"
     STABILITY_AI = "stability_ai"
     STABLE_DIFFUSION_API = "stable_diffusion_api"
+    COHERE = "cohere"

ops/txt_to_txt.py

@@ -0,0 +1,111 @@
+import asyncio
+import os
+import threading
+import queue
+from typing import Any, Callable, Optional
+
+from langchain import PromptTemplate, LLMChain
+from langchain.chat_models import ChatOpenAI
+from langchain.callbacks.base import BaseCallbackHandler
+from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
+from langchain.llms import Cohere
+from langchain.schema import (
+    AIMessage,
+    HumanMessage,
+    SystemMessage
+)
+
+from chisel.api.base_api_provider import BaseAPIProvider
+from chisel.data_types import Text, Image
+from chisel.ops.base_chisel import BaseChisel
+from chisel.ops.provider import Provider
+
+
+class TxtToTxt(BaseChisel):
+    def __init__(self, provider: Provider) -> None:
+        super().__init__(provider)
+        self.provider = provider
+
+    def _get_api(self, provider: Provider) -> BaseAPIProvider:
+        pass
+
+    def __call__(
+        self,
+        txt: Text,
+        callback: Optional[Callable] = None,
+    ) -> Text:
+        # TODO: generalize. This isn't going to scale.
+        if self.provider == "openai":
+            return openai_chat(txt, callback)
+        elif self.provider == "cohere":
+            return cohere_llm(txt, callback)
+
+
+class ThreadedGenerator:
+    def __init__(self):
+        self.queue = queue.Queue()
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        item = self.queue.get()
+        if item is StopIteration: raise item
+        return item
+
+    def send(self, data):
+        self.queue.put(data)
+
+    def close(self):
+        self.queue.put(StopIteration)
+
+
+class CustomCallback(BaseCallbackHandler):
+    def set_callback(self, callback):
+        self.callback = callback
+
+    def on_llm_new_token(self, token: str, **kwargs):
+        self.callback.on_llm_new_token(token, kwargs)
+
+    def on_llm_end(self, response, **kwargs: Any) -> Any:
+        self.callback.on_llm_end(response=response, kwargs=kwargs)
+
+
+def openai_chat_thread(g: ThreadedGenerator, prompt, callback):
+    custom_callback = CustomCallback()
+    custom_callback.set_callback(callback)
+    try:
+        chat = ChatOpenAI(
+            verbose=True,
+            streaming=True,
+            callbacks=[custom_callback],
+            temperature=0.7,
+        )
+        resp = chat([HumanMessage(content=prompt)])
+        return resp
+    finally:
+        g.close()
+
+
+def openai_chat(prompt, callback) -> ThreadedGenerator:
+    generator = ThreadedGenerator()
+    # threading.Thread(target=openai_chat_thread, args=(generator, prompt)).start()
+    resp = openai_chat_thread(generator, prompt, callback)
+    return resp
+
+
+def cohere_thread(prompt, callback):
+    custom_callback = CustomCallback()
+    custom_callback.set_callback(callback)
+
+    llm = Cohere(cohere_api_key=os.environ.get("COHERE_API_KEY"), temperature=0.7)
+
+    template = """{prompt}"""
+    prompt_template = PromptTemplate(template=template, input_variables=["prompt"])
+    llm_chain = LLMChain(llm=llm, prompt=prompt_template)
+    answer = llm_chain(prompt, callbacks=[custom_callback])
+
+
+def cohere_llm(prompt: Text, callback) -> None:
+    threading.Thread(target=cohere_thread, args=(prompt, callback)).start()
+    return None

pyproject.toml

@@ -17,7 +17,7 @@ google-cloud-storage = "^2.9.0"
 python = "^3.7"
 requests = "2.28.1"
 stability-sdk = "0.8.0"
-typer = {extras = ["all"], version = "^0.1.0"}
+typer = {extras = ["all"], version = "^0.9.0"}
 
 [build-system]
 requires = ["poetry-core"]