app.py

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

import gspread
from google.oauth2.service_account import Credentials

import os
import random
from huggingface_hub import login
import logging
import traceback

logging.basicConfig(level=logging.DEBUG)

HF_TOKEN = os.environ.get("HF_TOKEN", None)
login(token=HF_TOKEN)

MODEL = "5to9/hygge_v33-500_merged"

SPACER = "*"*50+"\n"

TITLE = """<h1><center>🇸🇪 Hygge – Möbelhauspoesie</center></h1>
<p><center><big><strong>Kann ein klitzekleines Sprachmodell eine komplexe Markensprache lernen?</strong> Diese Frage beantwortet Hygge mit einem fröhlichen &quot;Vielleicht!&quot; Das Modell produziert Texte im Stil eines fiktiven Möbelhauses. Auch wenn es mit 8B Parametern winzig ist, gibt es sich Mühe, so zu texten wie die Großen. Klappt schon ein bisschen. Mehr Infos <a href='https://portfolio.oliverheidorn.de/project/hygge-llm'>in meinem Portfolio.</a></big></center></p><p><center>Das Modell basiert auf <a href="https://huggingface.co/DiscoResearch/Llama3-DiscoLeo-Instruct-8B-v0.1">DiscoLeo 8B von Disco Research</a> einer deutschsprachig finegetuneten Version von <a href="https://huggingface.co/meta-llama/Meta-Llama-3-8B"> Llama 3</a>. Hygge speichert anonyme Logs der Ein- und Ausaben. Bitte benutz das Tool nur, wenn du damit cool bist. (Nicht wundern: Produkte haben die Namen von Simpsons-Charakteren. Is halt so.)</center></p></br>"""

PLACEHOLDER = """
<center>
<p>Lass uns lostexten!</p>
</center>
"""

SCOPES = ['https://www.googleapis.com/auth/spreadsheets', 'https://www.googleapis.com/auth/drive'] #spread scopes

examples = ['🛁 Schreibe einen Werbetext über ein platzsparendes Badmöbel.',
        '🚀 Schreibe einen Werbetext über ein gemütlich eingerichtetes Raumschiff für eine Marsreise.',
        '🪐 Schreibe einen Werbetext über die kuschelige Seite des Urknalls.',
        '🧠 Schreibe einen Werbetext, der Kants kategorischen Imperativ an einem Küchenhocker erklärt.',
        '🍌 Schreibe einen Werbetext über eine Küchenmöbelserie aus fermentierten Kumquats.',
        '🗄️ Schreibe einen Werbetext über unsere neuen stapelbaren Kleiderschränke für jede Wohnung.',
        '🪴 Schreibe einen Werbetext über einen praktischen Übertopf mit Handy-Ladevorrichtung.',
        '🤦‍♂️ Schreibe einen Werbetext, der unseren Beratungsservice bei der Einrichtung von Küchenschubladen auslobt.',
        '🍽️ Schreibe einen Werbetext über unseren Esstisch mit Wende-Tischplatte für gerade und ungerade Tage.',
        '🎄 Schreibe einen Werbetext über den Zauber von Weihnachten im Hochsommer.',
        '🕯️ Schreibe einen Werbetext, der begründet, warum man ein Möbelhaus UM HIMMELS WILLEN NIE ohne Teelichte verlassen sollte.',
        '😵‍💫 Schreib einen Werbetext über Fheisbqlezdkalxzsk.'
    ]


def generate_nested_examples(no_examples):
    return [[example] for example in random.sample(examples, no_examples)]

nested_examples = generate_nested_examples(12)


def starts_with_emoji(prompt_test):
    emojis = '🛁🚀🪐🧠🍌🗄️🪴🤦‍♂️🍽️🎄🕯️😵‍💫' 
    return  prompt_test[0] in emojis


def get_google_credentials():
    """Sets credentials for remote sheet"""
    service_account_info = {
        "type": "service_account",
        "project_id": os.environ.get("GOOGLE_PROJECT_ID"),
        "private_key_id": os.environ.get("GOOGLE_PRIVATE_KEY_ID"),
        "private_key": os.environ.get("GOOGLE_PRIVATE_KEY").replace('\\n', '\n'),
        "client_email": os.environ.get("GOOGLE_CLIENT_EMAIL"),
        "client_id": os.environ.get("GOOGLE_CLIENT_ID"),
        "auth_uri": os.environ.get("GOOGLE_AUTH_URI"),
        "token_uri": os.environ.get("GOOGLE_TOKEN_URI"),
        "auth_provider_x509_cert_url": os.environ.get("GOOGLE_AUTH_PROVIDER_CERT_URL"),
        "client_x509_cert_url": os.environ.get("GOOGLE_CLIENT_CERT_URL")
    }

    credentials = Credentials.from_service_account_info(service_account_info,scopes=SCOPES)
    return credentials


def get_google_sheet():
    """Intits auth, gets and returns instance of remote sheet"""
    credentials = get_google_credentials()
    client = gspread.authorize(credentials)
    sheet = client.open("hygge_inference_logs").sheet1  # Open your Google Sheet
    return sheet

device = "cuda" # set manually to reflect architecture of hf spaces

tokenizer = AutoTokenizer.from_pretrained(MODEL)
model = AutoModelForCausalLM.from_pretrained(
    MODEL,
    torch_dtype=torch.bfloat16,
    device_map="auto")

@spaces.GPU()
def stream_chat(message: str, 
                system_prompt: str,
                temperature: float = 0.8, 
                max_new_tokens: int = 300, 
                top_p: float = 1.0, 
                top_k: int = 50, 
                penalty: float = 1.2,
            ):

    system_prompt = """Du bist ein hochakkurater Chatbot, der humorvoll, einfach und natürlich formulieren kann und sich an Vorgaben hält wie ein guter Werbetexter. Du erhältst vom User eine Anweisung. Befolge sie genau, greife das Thema der Anweisung unbedingt auf und thematisiere es in deiner Antwort."""

    logging.debug(f'message: {message}, temperature: {temperature}, max_new_tokens: {max_new_tokens}')

    if starts_with_emoji(message):
        message = message[2:]

    conversation = [{"role": "system", "content": system_prompt},
                    {"role": "user", "content": message}]

    input_ids = tokenizer.apply_chat_template(conversation, add_generation_prompt=True, return_tensors="pt").to(model.device)
    #decoded_input = tokenizer.decode(input_ids[0], skip_special_tokens=False)
    #logging.debug(f'decoded_input: {decoded_input}')

    streamer = TextIteratorStreamer(tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=True)
    
    generate_kwargs = dict(
        input_ids=input_ids, 
        max_new_tokens = max_new_tokens,
        do_sample = False if temperature == 0 else True,
        top_p = top_p,
        top_k = top_k,
        temperature = temperature,
        repetition_penalty=penalty,
        eos_token_id=[128001,128008,128009],
        streamer=streamer,
    )

    with torch.no_grad():
        thread = Thread(target=model.generate, kwargs=generate_kwargs)
        thread.start()
        
    buffer = ""
    for new_text in streamer:
        buffer += new_text
        yield buffer
    
    try:
        sheet_row = [message, buffer, temperature, max_new_tokens, top_p, top_k, penalty]
        logging.debug(f'{SPACER}Output row: {sheet_row}\n{SPACER}') 
        sheet = get_google_sheet()
        sheet.append_row(sheet_row, table_range="A1:G1")
    except Exception as e:
        logging.error(f'{SPACER} Error: {e}, Traceback {traceback.format_exc()}')
    logging.debug(f'Generation done: {buffer}')



chatbot = gr.Chatbot(height=400, placeholder=PLACEHOLDER)

with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.HTML(TITLE)
    gr.ChatInterface(
        fn=stream_chat,
        chatbot=chatbot,
        fill_height=True,
        additional_inputs_accordion=gr.Accordion(label="⚙️ Parameters", open=False, render=False),
        additional_inputs=[
            gr.Slider(
                minimum=0,
                maximum=1,
                step=0.1,
                value=0.4,
                label="Temperature",
                render=False,
            ),
            gr.Slider(
                minimum=100,
                maximum=500,
                step=10,
                value=300,
                label="Max new tokens",
                render=False,
            ),
            gr.Slider(
                minimum=0.0,
                maximum=1.0,
                step=0.01,
                value=0.70,
                label="top_p",
                render=False,
            ),
            gr.Slider(
                minimum=1,
                maximum=50,
                step=1,
                value=10,
                label="top_k",
                render=False,
            ),
            gr.Slider(
                minimum=0.0,
                maximum=2.0,
                step=0.1,
                value=1.1,
                label="Repetition penalty",
                render=False,
            ),
        ],
        examples=nested_examples,
        cache_examples=False,
    )


if __name__ == "__main__":
    demo.launch()