Spaces:
Running on Zero
Running on Zero
| import re | |
| import gradio as gr | |
| import spaces | |
| import torch | |
| from threading import Thread | |
| from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer | |
| MODEL_ID = "CMSManhattan/JiRackDeltaNet_27b" | |
| tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True) | |
| model = AutoModelForCausalLM.from_pretrained( | |
| MODEL_ID, dtype=torch.bfloat16, trust_remote_code=True | |
| ) | |
| model.to("cuda") | |
| # Останавливаемся на конце реплики ассистента | |
| STOP_IDS = {tokenizer.eos_token_id} | |
| for tok in ("<|im_end|>", "<|endoftext|>"): | |
| tid = tokenizer.convert_tokens_to_ids(tok) | |
| if tid is not None and tid != tokenizer.unk_token_id: | |
| STOP_IDS.add(tid) | |
| STOP_IDS = [i for i in STOP_IDS if i is not None] | |
| THINK_RE = re.compile(r"<think>.*?</think>", re.DOTALL) | |
| def clean(text): | |
| text = THINK_RE.sub("", text) | |
| text = text.replace("<think>", "").replace("</think>", "") | |
| # если модель всё же начала выдуманную новую реплику, отрезаем её | |
| text = re.split(r"\n(?:user|assistant)\n", text)[0] | |
| return text.strip() | |
| def as_text(content): | |
| if isinstance(content, str): | |
| return content | |
| if isinstance(content, list): | |
| return "".join( | |
| p.get("text", "") if isinstance(p, dict) else str(p) for p in content | |
| ) | |
| return str(content) | |
| def build_messages(message, history): | |
| msgs = [] | |
| for turn in history: | |
| if isinstance(turn, dict): | |
| msgs.append({"role": turn["role"], "content": as_text(turn["content"])}) | |
| else: | |
| user, bot = turn | |
| msgs.append({"role": "user", "content": as_text(user)}) | |
| if bot: | |
| msgs.append({"role": "assistant", "content": as_text(bot)}) | |
| msgs.append({"role": "user", "content": as_text(message)}) | |
| return msgs | |
| def respond(message, history, max_new_tokens=512, temperature=0.7): | |
| inputs = tokenizer.apply_chat_template( | |
| build_messages(message, history), | |
| add_generation_prompt=True, | |
| return_tensors="pt", | |
| return_dict=True, | |
| ).to("cuda") | |
| streamer = TextIteratorStreamer( | |
| tokenizer, skip_prompt=True, skip_special_tokens=True | |
| ) | |
| Thread( | |
| target=model.generate, | |
| kwargs=dict( | |
| **inputs, | |
| streamer=streamer, | |
| max_new_tokens=int(max_new_tokens), | |
| do_sample=True, | |
| temperature=float(temperature), | |
| top_p=0.8, | |
| top_k=20, | |
| eos_token_id=STOP_IDS, | |
| ), | |
| ).start() | |
| out = "" | |
| for piece in streamer: | |
| out += piece | |
| yield clean(out) | |
| demo = gr.ChatInterface( | |
| respond, | |
| title="JiRack DeltaNet 27B", | |
| description="Try the JiRack DeltaNet 27B model.", | |
| additional_inputs=[ | |
| gr.Slider(64, 1024, value=512, step=64, label="Max new tokens"), | |
| gr.Slider(0.1, 1.5, value=0.7, step=0.1, label="Temperature"), | |
| ], | |
| ) | |
| demo.launch() |