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