use a local transformers model with a grammar

app.py

@@ -1,46 +1,63 @@
 import gradio as gr
-from huggingface_hub import InferenceClient
-import re
+import outlines
+import transformers
+import torch
 
 """
 For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
 """
-client = InferenceClient("HuggingFaceTB/SmolLM2-1.7B-Instruct")
+pipe = transformers.pipeline("text-generation", "HuggingFaceTB/SmolLM-135M-Instruct", torch_dtype=torch.float32)
+outlines_tokenizer = outlines.models.TransformerTokenizer(pipe.tokenizer)
+
+### TODO 1: use outliunes with a transformer model made directly
+### TODO 2: use a cfg
+
+def string_to_acrostic_grammar(s, dash_initial=True):
+    # this will convert a string to a CFG grammar
+    chars = filter(str.isalpha, s.upper())
+    grammar_rules = [('"- " ' if dash_initial else '') + f'"{char}" /[^-\\r\\n]+/ "\\n"' for char in chars]
+    return "?start: " + " ".join(grammar_rules)
+
+def is_this_prompt_a_list(prompt):
+    return False
+    # ask the model if the prompt is a list, by constraining the generation to yes or no about a question whether the prompt is a list
+    question = f'This is a prompt that you have been asked to answer:\n\n```\n{prompt}\n```\n\nIs this prompt asking for a list of items, instead of a story? Begin your answer with "Yes" if asking for a list, otherwise "No", and then give an explanation of why.'
+    grammar = '?start: ("Yes" | "No")'
+    cfg_logits_processor = outlines.processors.CFGLogitsProcessor(grammar, outlines_tokenizer)
+    output = pipe([{"role": "user", "content": question}], logits_processor=transformers.LogitsProcessorList([cfg_logits_processor]), max_new_tokens=10,)
+    # output = pipe([{"role": "system", "content": "You are a helpful assistant who answers in one-word answers."}, {"role": "user", "content": question}], max_new_tokens=10,)
+    response = output[0]['generated_text'][-1]['content']
+    print("is this prompt a list?", response)
+    return response == "Yes"
 
 
 def respond(
     message,
     history: list[tuple[str, str]],
     system_message,
+    acrostic,
     max_tokens,
     temperature,
     top_p,
 ):
-    messages = [{"role": "system", "content": system_message}]
-
-    for val in history:
-        if val[0]:
-            messages.append({"role": "user", "content": val[0]})
-        if val[1]:
-            messages.append({"role": "assistant", "content": val[1]})
-
-    messages.append({"role": "user", "content": message})
+    print({"message": message, "history": history, "system_message": system_message, "acrostic": acrostic, "max_tokens": max_tokens, "temperature": temperature, "top_p": top_p})
+    # grammar = "\n".join(['?start: item item item','?item: "- " /[^-\\r\\n]+/ "\\n"'])
+    grammar = string_to_acrostic_grammar(acrostic, dash_initial=is_this_prompt_a_list(message))
+    two_items_logits_processor = outlines.processors.CFGLogitsProcessor( grammar , outlines_tokenizer )
+    output = pipe([{"role": "user", "content": message}], logits_processor=transformers.LogitsProcessorList([two_items_logits_processor]), max_new_tokens=max_tokens,)
+    print(output)
+    response = output[0]['generated_text'][-1]['content']
 
-    response = ""
-    grammar_regex = r"^- H[^\r\n\x0b\x0c\x85\u2028\u2029]+\n- E[^\r\n\x0b\x0c\x85\u2028\u2029]+\n- L[^\r\n\x0b\x0c\x85\u2028\u2029]+\n- P[^\r\n\x0b\x0c\x85\u2028\u2029]+\n- M[^\r\n\x0b\x0c\x85\u2028\u2029]+\n- E[^\r\n\x0b\x0c\x85\u2028\u2029]+\n$"
+    # messages = [{"role": "system", "content": system_message}]
 
-    for message in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p,
-        response_format=grammar_regex,
-    ):
-        token = message.choices[0].delta.content
+    # for val in history:
+    #     if val[0]:
+    #         messages.append({"role": "user", "content": val[0]})
+    #     if val[1]:
+    #         messages.append({"role": "assistant", "content": val[1]})
 
-        response += token
-        yield response
+    # messages.append({"role": "user", "content": message})
+    yield response
 
 
 """
@@ -50,6 +67,7 @@ demo = gr.ChatInterface(
     respond,
     additional_inputs=[
         gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
+        gr.Textbox(value="I love you", label="acrostic"),
         gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Maximum new tokens"),
         gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
         gr.Slider(

requirements.txt

@@ -1 +1,5 @@
-huggingface_hub==0.25.2
+torch
+transformers
+outlines
+sentencepiece
+datasets