HedronCreeper/lfm2

1import gradio as gr2from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer3import torch4from threading import Thread5 6MODEL_NAMES = {7    "LFM 350M": "LiquidAI/LFM2-350M",8    "LFM 700M": "LiquidAI/LFM2-700M",9    "LFM 1.2B": "LiquidAI/LFM2-1.2B",10}11 12model_cache = {}13 14def load_model(model_key):15    if model_key in model_cache:16        return model_cache[model_key]17    model_name = MODEL_NAMES[model_key]18    tokenizer = AutoTokenizer.from_pretrained(model_name)19    device = "cuda" if torch.cuda.is_available() else "cpu"20    model = AutoModelForCausalLM.from_pretrained(21        model_name,22        dtype=torch.float16 if device == "cuda" else torch.float32,23    ).to(device)24    model_cache[model_key] = (tokenizer, model)25    return tokenizer, model26 27def chat_with_model(message, model_choice):28    tokenizer, model = load_model(model_choice)29    device = model.device30    31    # Absolute zero modification - your text goes straight to the AI32    prompt = message33    34    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)35    inputs = tokenizer(prompt, return_tensors="pt").to(device)36    37    generation_kwargs = dict(38        **inputs,39        streamer=streamer,40        max_new_tokens=1024,41        temperature=0.0,42        top_p=0.9,43        do_sample=True,44    )45    46    thread = Thread(target=model.generate, kwargs=generation_kwargs)47    thread.start()48    49    partial_text = ""50    for new_text in streamer:51        partial_text += new_text52        # Returns exactly one exchange: [User message, AI response]53        yield [[message, partial_text]]54 55def create_demo():56    # WhatsApp-inspired "Creeper" Dark Theme57    custom_theme = gr.themes.Soft(58        primary_hue="green",59        neutral_hue="slate",60    ).set(61        body_background_fill="*neutral_950",62        block_background_fill="*neutral_900",63        block_border_width="1px",64        block_label_text_color="*primary_500",65        button_primary_background_fill="*primary_600",66    )67 68    with gr.Blocks(theme=custom_theme, title="Creeper AI Chatbot") as demo:69        gr.Markdown("# 🌿 Creeper AI Chatbot")70        71        model_choice = gr.Dropdown(72            label="AI Brain (LFM)",73            choices=list(MODEL_NAMES.keys()),74            value="LFM 1.2B"75        )76        77        chatbot = gr.Chatbot(78            label="Chat View",79            height=500,80            bubble_full_width=False81        )82        83        with gr.Row():84            msg = gr.Textbox(85                label="Message",86                placeholder="Type here...",87                scale=4,88                show_label=False89            )90            submit_btn = gr.Button("Send", variant="primary", scale=1)91        92        clear = gr.Button("Clear Screen")93        94        # This handles the "No Memory" logic: 95        # Every time you hit send, it ignores history and just runs the current message.96        def start_chat(user_message):97            return "", [[user_message, None]]98 99        msg.submit(start_chat, [msg], [msg, chatbot]).then(100            chat_with_model, [msg, model_choice], chatbot101        )102        submit_btn.click(start_chat, [msg], [msg, chatbot]).then(103            chat_with_model, [msg, model_choice], chatbot104        )105        106        clear.click(lambda: None, None, chatbot, queue=False)107        108    return demo109 110if __name__ == "__main__":111    demo = create_demo()112    demo.queue()113    demo.launch(server_name="0.0.0.0", server_port=7860)114