File size: 2,795 Bytes
11b7390
 
 
 
 
7365dd9
11b7390
 
55f71ec
11b7390
e34fd28
 
 
 
11b7390
55f71ec
e34fd28
4e3d3cf
e34fd28
 
 
4e3d3cf
e34fd28
 
55f71ec
e34fd28
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
55f71ec
11b7390
55f71ec
2cb666f
55f71ec
11b7390
 
 
2cb666f
 
 
55f71ec
11b7390
 
2cb666f
 
11b7390
 
 
 
 
 
55f71ec
2cb666f
11b7390
55f71ec
 
11b7390
 
524e0e6
 
 
e34fd28
524e0e6
 
 
 
 
55f71ec
e34fd28
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
"""
Hugging Face Gradio Space for your style model.
Upload this as app.py to your HF Space.
"""

import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig
import torch

# ──────────────────────────────────────────────
MODELS = {
    "3B": "bugaku/igpt",
    "1.5B": "bugaku/igpt-smol",
}
# ──────────────────────────────────────────────

loaded_models = {}

def load_model(model_repo):
    if model_repo in loaded_models:
        return loaded_models[model_repo]

    print(f"Loading model: {model_repo}")
    tokenizer = AutoTokenizer.from_pretrained(model_repo)

    config = AutoConfig.from_pretrained(model_repo)
    if hasattr(config, "quantization_config"):
        delattr(config, "quantization_config")
    if hasattr(config, "_pre_quantization_dtype"):
        delattr(config, "_pre_quantization_dtype")

    model = AutoModelForCausalLM.from_pretrained(
        model_repo,
        config=config,
        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
        device_map="auto",
    )
    model.eval()
    loaded_models[model_repo] = (model, tokenizer)
    print(f"Loaded: {model_repo}")
    return model, tokenizer


def respond(message, history, model_choice, max_new_tokens, temperature):
    model_repo = MODELS[model_choice]
    model, tokenizer = load_model(model_repo)

    messages = [{"role": "user", "content": message}]

    encoded = tokenizer.apply_chat_template(
        messages,
        tokenize=True,
        add_generation_prompt=True,
        return_tensors="pt"
    )
    input_ids = encoded["input_ids"].to(model.device)
    attention_mask = encoded["attention_mask"].to(model.device)

    with torch.no_grad():
        outputs = model.generate(
            input_ids=input_ids,
            attention_mask=attention_mask,
            max_new_tokens=max_new_tokens,
            temperature=temperature,
            top_p=0.9,
            repetition_penalty=1.1,
            do_sample=True,
        )

    response_ids = outputs[0][input_ids.shape[-1]:]
    return tokenizer.decode(response_ids, skip_special_tokens=True)


with gr.Blocks(title="Style Model") as demo:
    gr.Markdown("## πŸ’¬ Style Model Chat")
    chatbot = gr.ChatInterface(
        fn=respond,
        additional_inputs=[
            gr.Dropdown(choices=list(MODELS.keys()), value=list(MODELS.keys())[0], label="Model"),
            gr.Slider(32, 512, value=256, step=32, label="Max response length"),
            gr.Slider(0.1, 2.0, value=0.8, step=0.1, label="Temperature"),
        ],
        chatbot=gr.Chatbot(height=500),
    )

demo.launch()