#!/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, torch_dtype=torch.bfloat16, device_map="auto") model = model.to("cuda" if torch.cuda.is_available() else "cpu") 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]: # Get Dutch date formatting locale try: locale.setlocale(locale.LC_ALL, "nl-NL") except locale.Error: pass now = datetime.now() now_str = now.strftime("%B %d, %Y, %H:%M:%S") # Reset locale try: locale.setlocale(locale.LC_ALL, locale.getdefaultlocale()) except locale.Error: pass conversation = [ { "role": "system", "content": f"Je bent 'Fietje' (of 'Fie' in het kort, naar 'Sofie'), een behulpzame en enthousiaste AI-assistent." f" Je werd gemaakt door Bram Vanroy, een onderzoeker aan de KU Leuven en het Instituut voor de Nederlandse Taal (INT)." f" De huidige datum en tijd is {now_str}." } ] 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"?"""], ["Maak een nieuw DnD personage met de naam 'Bram'. Geef een beschrijving, de skills, en de extra 'traits' met het aantal punten per vaardigheid. Gebruik JSON. Geef geen extra uitleg."], ["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"""\