#!/usr/bin/env python import os from threading import Thread from typing import Iterator import torch import gradio as gr # import spaces import torch from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer MAX_MAX_NEW_TOKENS = 2048 DEFAULT_MAX_NEW_TOKENS = 1024 MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "8192")) model_id = "BramVanroy/GEITje-7B-ultra" 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., do_sample: bool = False, ) -> Iterator[str]: conversation = [{"role": "system", "content": "Je bent 'GEITje Ultra', of 'Ultra' in het kort, een behulpzame en enthousiaste AI-assistent."}] 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, ) 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="GEITje-7B-ultra", show_share_button=True, avatar_images=(None, 'geitje-ultra-avatar.png')), 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.Checkbox( label="Do sample", value=False, ) ], examples=[ ["""Vraagje: welk woord hoort er niet in dit rijtje thuis: "auto, vliegtuig, geit, bus"?"""], ["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? Antwoord kort met enkel 'ja' of 'nee'."], ["Wat is groter, een olifant of het Atomium? 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?"], ], title="đ GEITje ultra đ¤", description="""\