app.py

import os
from threading import Thread
from typing import Iterator

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

# Konfigurationsparameter
MAX_MAX_NEW_TOKENS = 100
DEFAULT_MAX_NEW_TOKENS = 20
MAX_INPUT_TOKEN_LENGTH = 400  # Begrenzung auf 400 Tokens

# Modell und Tokenizer laden
model_id = "Loewolf/GPT_1"
model = AutoModelForCausalLM.from_pretrained(model_id)
tokenizer = AutoTokenizer.from_pretrained(model_id)

# Gradio Chat Interface Funktion
def generate(
    message: str,
    chat_history: list[tuple[str, str]],
    system_prompt: str,
    max_new_tokens: int = DEFAULT_MAX_NEW_TOKENS,
    temperature: float = 0.6,
    top_p: float = 0.9,
    top_k: int = 50,
    repetition_penalty: float = 1.2,
) -> str:
    conversation = []
    if system_prompt:
        conversation.append({"role": "system", "content": system_prompt})
    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(conversation, return_tensors="pt", truncation=True, max_length=MAX_INPUT_TOKEN_LENGTH)
    generate_kwargs = dict(
        input_ids=input_ids["input_ids"],
        max_length=input_ids["input_ids"].shape[1] + max_new_tokens,
        temperature=temperature,
        top_p=top_p,
        top_k=top_k,
        repetition_penalty=repetition_penalty,
        pad_token_id=tokenizer.eos_token_id
    )

    outputs = model.generate(**generate_kwargs)
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

# Gradio Interface
chat_interface = gr.Interface(
    fn=generate,
    inputs=[
        gr.Textbox(label="Message"),
        gr.JSON(label="Chat History"),
        gr.Textbox(label="System Prompt", lines=2),
        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.1, maximum=1.0, step=0.1, value=0.6),
        gr.Slider(label="Top-p (nucleus sampling)", minimum=0.05, maximum=1.0, step=0.05, value=0.9),
        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.2)
    ],
    outputs="text",
    live=True
)

# Starten des Gradio-Servers
if __name__ == "__main__":
    chat_interface.launch()