Use text-generation inference 🔥🔥

app.py

@@ -1,12 +1,16 @@
+import json
 import os
-from threading import Thread
 
 import gradio as gr
-import torch
-from transformers import (AutoModelForCausalLM, AutoTokenizer,
-                          TextIteratorStreamer, set_seed)
+# import torch
+# from transformers import (AutoModelForCausalLM, AutoTokenizer,
+#                           TextIteratorStreamer, set_seed)
 from huggingface_hub import Repository
-import json
+from text_generation import Client
+
+# from threading import Thread
+
+
 
 theme = gr.themes.Monochrome(
     primary_hue="indigo",
@@ -16,27 +20,32 @@ theme = gr.themes.Monochrome(
     font=[gr.themes.GoogleFont("Open Sans"), "ui-sans-serif", "system-ui", "sans-serif"],
 )
 HF_TOKEN = os.environ.get("HF_TOKEN", None)
-os.environ["TOKENIZERS_PARALLELISM"] = "false"
+# os.environ["TOKENIZERS_PARALLELISM"] = "false"
 if HF_TOKEN:
     repo = Repository(
         local_dir="data", clone_from="trl-lib/stack-llama-prompts", use_auth_token=HF_TOKEN, repo_type="dataset"
     )
 
+client = Client(
+    "https://api-inference.huggingface.co/models/trl-lib/llama-se-rl-merged",
+    headers={"Authorization": f"Bearer {HF_TOKEN}"},
+)
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model_id = "trl-lib/llama-se-rl-merged"
-print(f"Loading model: {model_id}")
-if device == "cpu":
-    model = AutoModelForCausalLM.from_pretrained(model_id, low_cpu_mem_usage=True, use_auth_token=HF_TOKEN)
-else:
-    model = AutoModelForCausalLM.from_pretrained(
-        model_id, device_map="auto", load_in_8bit=True, use_auth_token=HF_TOKEN
-    )
+# device = "cuda" if torch.cuda.is_available() else "cpu"
+# model_id = "trl-lib/llama-se-rl-merged"
+# print(f"Loading model: {model_id}")
+# if device == "cpu":
+#     model = AutoModelForCausalLM.from_pretrained(model_id, low_cpu_mem_usage=True, use_auth_token=HF_TOKEN)
+# else:
+#     model = AutoModelForCausalLM.from_pretrained(
+#         model_id, device_map="auto", load_in_8bit=True, use_auth_token=HF_TOKEN
+#     )
 
-tokenizer = AutoTokenizer.from_pretrained(model_id, use_auth_token=HF_TOKEN)
+# tokenizer = AutoTokenizer.from_pretrained(model_id, use_auth_token=HF_TOKEN)
 
 PROMPT_TEMPLATE = """Question: {prompt}\n\nAnswer:"""
 
+
 def save_inputs_and_outputs(inputs, outputs, generate_kwargs):
     with open(os.path.join("data", "prompts.jsonl"), "a") as f:
         json.dump({"inputs": inputs, "outputs": outputs, "generate_kwargs": generate_kwargs}, f, ensure_ascii=False)
@@ -44,44 +53,101 @@ def save_inputs_and_outputs(inputs, outputs, generate_kwargs):
         commit_url = repo.push_to_hub()
 
 
-def generate(instruction, temperature=0.9, max_new_tokens=128, top_p=0.95, top_k=100):
-    set_seed(42)
+# def generate(instruction, temperature=0.9, max_new_tokens=128, top_p=0.95, top_k=100):
+#     set_seed(42)
+#     formatted_instruction = PROMPT_TEMPLATE.format(prompt=instruction)
+
+#     temperature = float(temperature)
+#     top_p = float(top_p)
+#     streamer = TextIteratorStreamer(tokenizer)
+#     model_inputs = tokenizer(formatted_instruction, return_tensors="pt", truncation=True, max_length=2048).to(device)
+
+#     generate_kwargs = dict(
+#         top_p=top_p,
+#         temperature=temperature,
+#         max_new_tokens=max_new_tokens,
+#         do_sample=True,
+#         top_k=top_k,
+#         eos_token_id=tokenizer.eos_token_id,
+#         pad_token_id=tokenizer.eos_token_id,
+#     )
+#     t = Thread(target=model.generate, kwargs={**dict(model_inputs, streamer=streamer), **generate_kwargs})
+#     t.start()
+
+#     output = ""
+#     hidden_output = ""
+#     for new_text in streamer:
+#         # skip streaming until new text is available
+#         if len(hidden_output) <= len(formatted_instruction):
+#             hidden_output += new_text
+#             continue
+#         # replace eos token
+#         # if tokenizer.eos_token in new_text:
+#         #     new_text = new_text.replace(tokenizer.eos_token, "")
+#         output += new_text
+#         yield output
+#     if HF_TOKEN:
+#         print("Pushing prompt and completion to the Hub")
+#         save_inputs_and_outputs(formatted_instruction, output, generate_kwargs)
+#     return output
+
+
+def generate(instruction, temperature=0.9, max_new_tokens=256, top_p=0.95, top_k=100):
+    # set_seed(42)
     formatted_instruction = PROMPT_TEMPLATE.format(prompt=instruction)
 
     temperature = float(temperature)
     top_p = float(top_p)
-    streamer = TextIteratorStreamer(tokenizer)
-    model_inputs = tokenizer(formatted_instruction, return_tensors="pt", truncation=True, max_length=2048).to(device)
 
-    generate_kwargs = dict(
-        top_p=top_p,
+    stream = client.generate_stream(
+        formatted_instruction,
         temperature=temperature,
+        truncate=999,
         max_new_tokens=max_new_tokens,
-        do_sample=True,
+        top_p=top_p,
         top_k=top_k,
-        eos_token_id=tokenizer.eos_token_id,
-        pad_token_id=tokenizer.eos_token_id,
+        # stop_sequences=["</s>"],
     )
-    t = Thread(target=model.generate, kwargs={**dict(model_inputs, streamer=streamer), **generate_kwargs})
-    t.start()
 
     output = ""
-    hidden_output = ""
-    for new_text in streamer:
-        # skip streaming until new text is available
-        if len(hidden_output) <= len(formatted_instruction):
-            hidden_output += new_text
-            continue
-        # replace eos token
-        # if tokenizer.eos_token in new_text:
-        #     new_text = new_text.replace(tokenizer.eos_token, "")
-        output += new_text
+    for response in stream:
+        output += response.token.text
         yield output
-    if HF_TOKEN:
-        print("Pushing prompt and completion to the Hub")
-        save_inputs_and_outputs(formatted_instruction, output, generate_kwargs)
+
     return output
 
+    # streamer = TextIteratorStreamer(tokenizer)
+    # model_inputs = tokenizer(formatted_instruction, return_tensors="pt", truncation=True, max_length=2048).to(device)
+
+    # generate_kwargs = dict(
+    #     top_p=top_p,
+    #     temperature=temperature,
+    #     max_new_tokens=max_new_tokens,
+    #     do_sample=True,
+    #     top_k=top_k,
+    #     # eos_token_id=tokenizer.eos_token_id,
+    #     # pad_token_id=tokenizer.eos_token_id,
+    # )
+    # t = Thread(target=model.generate, kwargs={**dict(model_inputs, streamer=streamer), **generate_kwargs})
+    # t.start()
+
+    # output = ""
+    # hidden_output = ""
+    # for new_text in streamer:
+    #     # skip streaming until new text is available
+    #     if len(hidden_output) <= len(formatted_instruction):
+    #         hidden_output += new_text
+    #         continue
+    #     # replace eos token
+    #     # if tokenizer.eos_token in new_text:
+    #     #     new_text = new_text.replace(tokenizer.eos_token, "")
+    #     output += new_text
+    #     yield output
+    # if HF_TOKEN:
+    #     print("Pushing prompt and completion to the Hub")
+    #     save_inputs_and_outputs(formatted_instruction, output, generate_kwargs)
+    # return output
+
 
 examples = [
     "A llama is in my lawn. How do I get rid of him?",
@@ -167,4 +233,4 @@ with gr.Blocks(theme=theme, analytics_enabled=False, css=".generating {visibilit
     instruction.submit(generate, inputs=[instruction, temperature, max_new_tokens, top_p, top_k], outputs=[output])
 
 demo.queue(concurrency_count=1)
-demo.launch(enable_queue=True)#, share=True)
+demo.launch(enable_queue=True)  # , share=True)

requirements.txt

@@ -1,7 +1,7 @@
 huggingface_hub
-bitsandbytes
-sentencepiece
-git+https://github.com/huggingface/transformers.git@98268b2e76189d65f7068625cf382ebe03b98480
-accelerate>=0.16.0
-bitsandbytes
-sentencepiece
+# bitsandbytes
+# sentencepiece
+# git+https://github.com/huggingface/transformers.git@98268b2e76189d65f7068625cf382ebe03b98480
+# accelerate>=0.16.0
+# bitsandbytes
+# sentencepiece