Initial commit

app.py

@@ -0,0 +1,107 @@
+import gradio as gr
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from backtrack_sampler import BacktrackSampler, CreativeWritingStrategy
+from backtrack_sampler.provider.transformers_provider import TransformersProvider
+import torch
+import asyncio
+
+description = """## Compare Creative Writing: Custom Sampler vs. Backtrack Sampler with Creative Writing Strategy
+This is a demo of [Backtrack Sampler](https://github.com/Mihaiii/backtrack_sampler) using one of its algorithms named "Creative Writing Strategy".
+<br />On the left you have the output of the standard sampling and on the write the output privided by Backtrack Sampler.
+"""
+# Load tokenizer
+model_name = "unsloth/Llama-3.2-1B-Instruct"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+# Load two instances of the model on CUDA for parallel inference
+model1 = AutoModelForCausalLM.from_pretrained(model_name).to("cuda")
+
+model2 = AutoModelForCausalLM.from_pretrained(model_name)
+device = torch.device('cuda')
+
+strategy = CreativeWritingStrategy(top_p_flat = 0.8, top_k_threshold_flat = 2, min_prob_second_highest = 0.2)
+provider = TransformersProvider(model2, tokenizer, device)
+creative_sampler = BacktrackSampler(strategy, provider)
+
+# Helper function to create message array for the chat template
+def create_chat_template_messages(history, prompt):
+    messages = [{"role": "user", "content": prompt}]
+    
+    for i, (input_text, response_text) in enumerate(history):
+        messages.append({"role": "user" if i % 2 == 0 else "assistant", "content": input_text})
+        messages.append({"role": "assistant", "content": response_text})
+    
+    return messages
+
+# Async function for generating responses using two models
+@spaces.GPU(duration=60)
+async def generate_responses(prompt, history):
+    # Create messages array for chat history and apply template
+    messages = create_chat_template_messages(history, prompt)
+    wrapped_prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_special_tokens=True, add_generation_prompt=True)
+
+    #already has special tokens
+    inputs = tokenizer.encode(wrapped_prompt, add_special_tokens=False, return_tensors="pt").to("cuda")
+    # Standard sampler task
+    standard_task = asyncio.to_thread(
+        model1.generate, inputs, max_length=2048, temperature=1
+    )
+
+    # Custom sampler task: loop over generator and collect outputs in a list
+    async def custom_sampler_task():
+        generated_list = []
+        generator = creative_sampler.generate(wrapped_prompt, max_length=2048, temperature=1)
+        for token in generator:
+            generated_list.append(token)
+        return tokenizer.decode(generated_list, skip_special_tokens=True)
+
+    # Wait for both responses
+    standard_output, custom_output = await asyncio.gather(standard_task, custom_sampler_task())
+    # Decode standard output and remove the prompt from the generated response
+    standard_response = tokenizer.decode(standard_output[0][len(inputs[0]):], skip_special_tokens=True)
+
+    return standard_response.strip(), custom_output.strip()
+
+# Create the Gradio interface with the Citrus theme
+with gr.Blocks(theme=gr.themes.Citrus()) as demo:
+    gr.Markdown(description)
+
+    # Chatbot components
+    with gr.Row():
+        standard_chat = gr.Chatbot(label="Standard Sampler")
+        custom_chat = gr.Chatbot(label="Creative Writing Strategy")
+
+    # Input components
+    with gr.Row():
+        prompt_input = gr.Textbox(label="Enter your prompt", placeholder="Type your message here...", lines=1)
+
+    # Example prompts
+    examples = [
+        "Write me a short story about a talking dog who wants to be a detective.",
+        "Tell me a short tale of a dragon who is afraid of heights.",
+        "Create a short story where aliens land on Earth, but they just want to throw a party."
+    ]
+
+    # Add example buttons
+    gr.Examples(examples=examples, inputs=prompt_input)
+
+    # Button to submit the prompt
+    submit_button = gr.Button("Submit")
+
+    # Function to handle chat updates
+    async def update_chat(prompt, standard_history, custom_history):
+        standard_response, custom_response = await generate_responses(prompt, standard_history)
+
+        # Append new responses to chat histories
+        standard_history = standard_history + [(prompt, standard_response)]
+        custom_history = custom_history + [(prompt, custom_response)]
+
+        # Clear the input field after submission
+        return standard_history, custom_history, ""
+
+    # Bind the submit button to the update function and allow pressing Enter to submit
+    prompt_input.submit(fn=update_chat, inputs=[prompt_input, standard_chat, custom_chat], outputs=[standard_chat, custom_chat, prompt_input])
+    submit_button.click(fn=update_chat, inputs=[prompt_input, standard_chat, custom_chat], outputs=[standard_chat, custom_chat, prompt_input])
+
+# Launch the app with queueing and sharing enabled
+demo.queue().launch(share=True, debug=True)

requirements.txt

@@ -0,0 +1 @@
+backtrack_sampler