| import spaces | |
| import torch | |
| import gradio as gr | |
| from huggingface_hub import hf_hub_download | |
| from safetensors.torch import load_file | |
| MODEL_REPO = "ClokAI/ci-base" | |
| print("Loading model...") | |
| from clokai import CiModel, CiConfig | |
| from transformers import AutoTokenizer | |
| config = CiConfig() | |
| model = CiModel(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"<s> 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="ClokAI ci-base Test", | |
| description="Test ci-base model - safetensors + clokai library", | |
| 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"), | |
| ], | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch() | |