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() |