| import threading |
|
|
| import spaces |
| import torch |
| import gradio as gr |
| from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer |
|
|
| MODEL_ID = "KreatixLabs/Kreatix-K1" |
| DEFAULT_SYSTEM = "You are Kreatix K1, created by KreatixLabs. You are a helpful coding assistant." |
|
|
| tokenizer = AutoTokenizer.from_pretrained(MODEL_ID) |
| model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype=torch.float16).to("cuda") |
| model.eval() |
|
|
|
|
| def to_text(content): |
| if isinstance(content, str): |
| return content |
| if isinstance(content, dict): |
| return content.get("text", "") or "" |
| if isinstance(content, (list, tuple)): |
| return "".join(to_text(p) for p in content) |
| return "" if content is None else str(content) |
|
|
|
|
| @spaces.GPU(duration=120) |
| def chat(message, history, system_prompt, temperature, top_p, max_tokens): |
| messages = [] |
| if system_prompt and system_prompt.strip(): |
| messages.append({"role": "system", "content": system_prompt.strip()}) |
| for m in history or []: |
| role = m.get("role") if isinstance(m, dict) else None |
| if role in ("user", "assistant"): |
| text = to_text(m.get("content")) |
| if text: |
| messages.append({"role": role, "content": text}) |
| messages.append({"role": "user", "content": to_text(message)}) |
|
|
| prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) |
| inputs = tokenizer(prompt, return_tensors="pt").to("cuda") |
| streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True) |
| temperature = float(temperature) |
| kwargs = dict( |
| **inputs, |
| max_new_tokens=int(max_tokens), |
| do_sample=temperature > 0.05, |
| temperature=max(temperature, 0.05), |
| top_p=float(top_p), |
| repetition_penalty=1.05, |
| pad_token_id=tokenizer.eos_token_id, |
| streamer=streamer, |
| ) |
| worker = threading.Thread(target=model.generate, kwargs=kwargs) |
| worker.start() |
| out = "" |
| for tok in streamer: |
| out += tok |
| yield out |
| worker.join() |
|
|
|
|
| demo = gr.ChatInterface( |
| fn=chat, |
| title="Kreatix K1", |
| description="A 7B open coding model by KreatixLabs. Weights and GGUF: [KreatixLabs/Kreatix-K1](https://huggingface.co/KreatixLabs/Kreatix-K1)", |
| additional_inputs=[ |
| gr.Textbox(value=DEFAULT_SYSTEM, label="System prompt", lines=2), |
| gr.Slider(0.0, 2.0, value=0.7, step=0.05, label="Temperature"), |
| gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="Top-p"), |
| gr.Slider(64, 2048, value=768, step=64, label="Max new tokens"), |
| ], |
| examples=[ |
| ["Who are you?"], |
| ["Write a Python function that merges two sorted lists."], |
| ["My code won't work."], |
| ["Explain async/await in JavaScript like I'm tired."], |
| ], |
| cache_examples=False, |
| fill_height=True, |
| ) |
|
|
| demo.queue(default_concurrency_limit=4).launch() |
|
|
