🍻 cheers

app.py

@@ -0,0 +1,28 @@
+from threading import Thread
+
+from closedai import ClosedAIPipeline
+from closedai.server import app, data  # noqa
+
+from transformers import LlamaForCausalLM, LlamaTokenizer, TextIteratorStreamer, pipeline
+
+class LlamaPipeline(ClosedAIPipeline):
+    def __init__(self):
+        tokenizer = LlamaTokenizer.from_pretrained("decapoda-research/llama-7b-hf")
+        model = LlamaForCausalLM.from_pretrained("decapoda-research/llama-7b-hf")
+        self.streamer = TextIteratorStreamer(tokenizer, skip_prompt=True)
+        self.pipe = pipeline(
+            'text-generation',
+            model=model,
+            tokenizer=tokenizer,
+            streamer=self.streamer,
+            device="auto"
+        )
+
+    def generate_completion(self, text, **generate_kwargs):
+        thread = Thread(target=self.pipe.__call__, kwargs=dict(text_inputs=text, **generate_kwargs))
+        thread.start()
+        for new_text in self.streamer:
+            yield new_text
+
+pipeline = LlamaPipeline()
+data["pipeline"] = pipeline

requirements.txt

@@ -0,0 +1,4 @@
+git+https://github.com/nateraw/closedai.git
+git+https://github.com/huggingface/transformers.git
+sentencepiece
+accelerate