import gradio as gr import torch from huggingface_hub import hf_hub_download from safetensors.torch import load_file MODEL_REPO = "clokai/CLOK-CEM" print("Loading model...") from configuration_clokcem import ClokcemConfig from modeling_clokcem import ClokcemForCausalLM from transformers import AutoTokenizer config = ClokcemConfig() model = ClokcemForCausalLM(config) weights_path = hf_hub_download(repo_id=MODEL_REPO, filename="model.safetensors") sd = load_file(weights_path) model.load_state_dict(sd, strict=False) model.eval() tokenizer = AutoTokenizer.from_pretrained(MODEL_REPO) print("Model ready!") def chat(message, history, temperature, max_tokens): device = next(model.parameters()).device formatted = f"<|system|>You are a helpful customer care assistant.<|user|>{message}<|assistant|>" inputs = tokenizer(formatted, return_tensors="pt").to(device) generated = [] input_ids = inputs["input_ids"] for _ in range(max_tokens): with torch.no_grad(): logits = model(input_ids).logits probs = torch.softmax(logits[:, -1, :] / temperature, dim=-1) next_token = torch.multinomial(probs, num_samples=1) if next_token.item() == tokenizer.eos_token_id: break generated.append(next_token.item()) input_ids = torch.cat([input_ids, next_token], dim=-1) return tokenizer.decode(generated, skip_special_tokens=True) demo = gr.ChatInterface( fn=chat, title="ClokCEM - Customer Executive Model", description="354M parameter model for enterprise customer support", additional_inputs=[ gr.Slider(0.1, 2.0, value=0.6, step=0.1, label="Temperature"), gr.Slider(50, 500, value=200, step=50, label="Max Tokens"), ], ) demo.launch() torch transformers safetensors huggingface_hub gradio