app.py

#!/usr/bin/env python

from datetime import datetime
import locale
import os
from threading import Thread
from typing import Iterator

import gradio as gr
import spaces
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer

MAX_MAX_NEW_TOKENS = 1536
DEFAULT_MAX_NEW_TOKENS = 1024
MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "8192"))

model_id = "BramVanroy/fietje-2b-chat"
avatar_url = "https://huggingface.co/spaces/BramVanroy/fietje-2b/resolve/main/img/fietje-2b-avatar.png"
model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", load_in_8bit=True)
tokenizer = AutoTokenizer.from_pretrained(model_id)
tokenizer.pad_token_id = tokenizer.eos_token_id


@spaces.GPU
def generate(
    message: str,
    chat_history: list[tuple[str, str]],
    max_new_tokens: int = 1024,
    temperature: float = 1,
    top_p: float = 1.,
    top_k: int = 50,
    repetition_penalty: float = 1.,
    no_repeat_ngram_size: int = 4,
    do_sample: bool = False,
) -> Iterator[str]:
    conversation = []

    for user, assistant in chat_history:
        conversation.extend([{"role": "user", "content": user}, {"role": "assistant", "content": assistant}])
    conversation.append({"role": "user", "content": message})

    input_ids = tokenizer.apply_chat_template(conversation, add_generation_prompt=True, return_tensors="pt")
    if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH:
        input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]
        gr.Warning(f"Trimmed input from conversation as it was longer than {MAX_INPUT_TOKEN_LENGTH} tokens.")
    input_ids = input_ids.to(model.device)

    streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True)
    generate_kwargs = dict(
        {"input_ids": input_ids},
        streamer=streamer,
        max_new_tokens=max_new_tokens,
        do_sample=do_sample,
        top_p=top_p,
        top_k=top_k,
        temperature=temperature,
        num_beams=1,
        repetition_penalty=repetition_penalty,
        no_repeat_ngram_size=no_repeat_ngram_size,
    )
    t = Thread(target=model.generate, kwargs=generate_kwargs)
    t.start()

    outputs = []
    for text in streamer:
        outputs.append(text)
        yield "".join(outputs)


chat_interface = gr.ChatInterface(
    fn=generate,
    chatbot=gr.Chatbot(height=450,
                      label="Fietje",
                      show_share_button=True,
                      avatar_images=(None, avatar_url)),
    additional_inputs=[
        gr.Slider(
            label="Max new tokens",
            minimum=1,
            maximum=MAX_MAX_NEW_TOKENS,
            step=1,
            value=DEFAULT_MAX_NEW_TOKENS,
        ),
        gr.Slider(
            label="Temperature",
            minimum=0.05,
            maximum=2,
            step=0.05,
            value=1.0,
        ),
        gr.Slider(
            label="Top-p (nucleus sampling)",
            minimum=0.05,
            maximum=1.0,
            step=0.05,
            value=1.0,
        ),
        gr.Slider(
            label="Top-k",
            minimum=1,
            maximum=1000,
            step=1,
            value=50,
        ),
        gr.Slider(
            label="Repetition penalty",
            minimum=1.0,
            maximum=2.0,
            step=0.05,
            value=1.,
        ),
        gr.Slider(
            label="No repeat n-gram",
            minimum=0,
            maximum=20,
            step=1,
            value=0,
        ),
        gr.Checkbox(
            label="Do sample",
            value=False,
        )
    ],
    examples=[
        ["""Vraagje: welk woord hoort er niet in dit rijtje thuis: "auto, vliegtuig, geit, bus"?"""],
        ["Wat is het grootste in de winter, een Aziatische olifant of het Atomium?"],
        ["Schrijf een nieuwsbericht voor De Speld over de inzet van een kudde geiten door het Nederlands Forensisch Instituut"],
        ["Wat zijn drie leuke dingen om te doen als ik een weekendje naar Belgisch Limburg ga?"],
        ["Met wie trad clown Bassie op?"],
        ["Kan je naar de maan fietsen? Redeneer stap voor stap."],
        ["Ik wil in Python een nieuwe abstracte klasse aanmaken die `Dier` heet en die een abstracte methode `geluid_maken` heeft. Maak daarnaast ook een subklasse `Geit` aan, met een passende invulling van `geluid_maken` voor dit mekkerende dier."],
        ["Wat is het belang van open-source taalmodellen?"],
    ],
    cache_examples=False,
    title="Fietje",
    description=f"""\
<div style="float: left; margin-right: 1em; margin-bottom: 0.48em; max-width: 96px;">
    <img src={avatar_url} style="display: inline-block; width: auto; height: auto;border-radius:50%">
    <div style="clear: both"></div>
</div>

[Fietje](https://huggingface.co/BramVanroy/fietje-2b-chat) is een open en efficiënte chatbot, gebaseerd op phi-2 van Microsoft. Hoewel het model slechts 2,7 miljard parameters heeft, scoort het [in benchmarks](https://github.com/BramVanroy/fietje?tab=readme-ov-file#performance) bijna zo goed als, en soms zelfs beter dan!, modellen die meer dan dubbel zo groot zijn, zoals [GEITje 7B Ultra](https://huggingface.co/BramVanroy/GEITje-7B-ultra). In de praktijk zal Fietje echter een minder goede assistent zijn, wat vaker hallucineren, en minder kunnen en weten dan GEITje 7B Ultra. Je zal haar moeten vergeven! Ze is dan wel een stuk kleiner dan GEITje, en niet zo krachtig, maar wel veel sneller! Kleine, krachtige modellen zijn nuttig om LLMs toegankelijker te maken zodat je hen eenvoudiger op je eigen apparaten, tot zelfs op je telefoon!, kan gebruiken.

Meer informatie vind je terug op [de Github-pagina](https://github.com/BramVanroy/fietje).

""",
    submit_btn="Genereer",
    stop_btn="Stop",
    retry_btn="🔄 Opnieuw",
    undo_btn="↩️ Ongedaan maken",
    clear_btn="🗑️ Wissen",
)

with gr.Blocks(css="style.css") as demo:
    chat_interface.render()
    gr.Markdown("""\
## Opmerkingen

Hoewel dit model gealigneerd is met AI feedback (van gpt-4-turbo), kan en zal het nog steeds fouten maken, leugens verzinnen, redeneringsfouten maken, en misschien wel stoute dingen vertellen. Gebruik dit model dus op eigen risico en controleer de output zelf!

Als je problemen ondervindt, [rapporteer die dan gerust](https://huggingface.co/spaces/BramVanroy/fietje-2b/discussions), al zal ik er waarschijnlijk niet veel aan kunnen verhelpen! Toch is het goed om een verslag bij te houden van mogelijke vooroordelen of problemen die het model heeft, zodat daar in toekomstige ontwikkelingen rekening mee gehouden kan worden. 

Mijn eerste gevoel is dat de chat-variant wel eens kan verdwalen in lange opsommingen en ook onverwachts willekeurige karakters kan uitspugen. Mijn vermoeden is dat het DPO-proces niet helemaal goed gelopen is.
""")

if __name__ == "__main__":
    demo.queue(max_size=20).launch()