first commit

app.py

@@ -0,0 +1,144 @@
+import gradio as gr
+from threading import Thread
+from open_lm.hf import *
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
+import torch
+from gradio.layouts import Accordion
+
+# Define model options
+MODEL_OPTIONS = {
+    "TRI-ML/DCLM-1B": "TRI-ML/DCLM-1B",
+    "Apple DCLM-Baseline-7B": "apple/DCLM-Baseline-7B"
+}
+
+# Global variables for model and tokenizer
+current_model = None
+current_tokenizer = None
+
+def load_model(model_name):
+    global current_model, current_tokenizer
+    current_tokenizer = AutoTokenizer.from_pretrained(MODEL_OPTIONS[model_name])
+    current_model = AutoModelForCausalLM.from_pretrained(MODEL_OPTIONS[model_name])
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    current_model = current_model.to(device)
+    return f"Loaded model: {model_name}"
+
+def generate(
+    prompt, model_choice, temperature=0.9, max_new_tokens=256, top_p=0.95, repetition_penalty=1.0,
+):
+    global current_model, current_tokenizer
+    
+    if current_model is None or current_tokenizer is None:
+        return "Please load a model first."
+
+    temperature = float(temperature)
+    if temperature < 1e-2:
+        temperature = 1e-2
+    top_p = float(top_p)
+    
+    inputs = current_tokenizer(prompt, return_tensors="pt").to(current_model.device)
+    
+    generate_kwargs = dict(
+        **inputs,
+        max_new_tokens=max_new_tokens,
+        temperature=temperature,
+        top_p=top_p,
+        repetition_penalty=repetition_penalty,
+        do_sample=True,
+        pad_token_id=current_tokenizer.eos_token_id
+    )
+
+    streamer = TextIteratorStreamer(current_tokenizer, skip_prompt=True, skip_special_tokens=True)
+    generate_kwargs["streamer"] = streamer
+
+    thread = Thread(target=current_model.generate, kwargs=generate_kwargs)
+    thread.start()
+
+    # Write the prompt in blue
+    output = "<span style='color: blue;'>" + prompt + "</span>"
+    for new_text in streamer:
+        if isinstance(new_text, torch.Tensor):
+            new_text = current_tokenizer.decode(new_text)
+        output += new_text
+        yield output
+
+    thread.join()
+    return output
+
+additional_inputs=[
+        gr.Slider(
+            label="Temperature",
+            value=0.9,
+            minimum=0.0,
+            maximum=1.0,
+            step=0.05,
+            interactive=True,
+            info="Higher values produce more diverse outputs",
+        ),
+        gr.Slider(
+            label="Max new tokens",
+            value=256,
+            minimum=0,
+            maximum=1048,
+            step=64,
+            interactive=True,
+            info="The maximum numbers of new tokens",
+        ),
+        gr.Slider(
+            label="Top-p (nucleus sampling)",
+            value=0.90,
+            minimum=0.0,
+            maximum=1,
+            step=0.05,
+            interactive=True,
+            info="Higher values sample more low-probability tokens",
+        ),
+        gr.Slider(
+            label="Repetition penalty",
+            value=1.2,
+            minimum=1.0,
+            maximum=2.0,
+            step=0.05,
+            interactive=True,
+            info="Penalize repeated tokens",
+        )
+    ]
+
+with gr.Blocks() as demo:
+    gr.Markdown(
+        """
+        # DCLM Text Completion Demo
+        This demo allows you to generate text using a DCLM model. 
+        These models are trained to predict the next word in a sequence of text, and can be used to generate text completions, they are not chatbots. 
+        
+        First select a model from the dropdown and click "Load Model". 
+        Then enter some text in the text box and click "Generate" to see the model's completion.
+        """
+    )
+
+    
+    with gr.Row():
+        model_dropdown = gr.Dropdown(choices=list(MODEL_OPTIONS.keys()), label="Select Model")
+
+    model_dropdown.select(
+        load_model,
+        inputs=[model_dropdown],
+        outputs=[gr.Textbox(label="Model Status")]
+    )
+
+    text_input = gr.Textbox(lines=3, label="Input Text")
+    text_output = gr.HTML(label="Generated Text")
+
+    generate_button = gr.Button("Generate")
+
+    generate_button.click(
+        generate,
+        inputs=[text_input, model_dropdown, *additional_inputs],
+        outputs=[text_output]
+    )
+    with Accordion(label="Advanced Options", open=False):
+        for input_component in additional_inputs:
+            if not input_component.is_rendered:
+                input_component.render()
+
+demo.launch()

requirements.txt

@@ -0,0 +1,4 @@
+git+https://github.com/mlfoundations/open_lm.git
+gradio
+transformers
+torch