| """Pinned text-only inference. Model weights are loaded only on a cloud GPU.""" |
| from __future__ import annotations |
|
|
| import json |
| import os |
| from pathlib import Path |
|
|
| ROOT = Path(__file__).resolve().parent.parent |
| CONFIG = json.loads(Path(__file__).with_name("config.json").read_text(encoding="utf-8")) |
| os.environ.setdefault("HF_HOME", str(ROOT / ".cache/huggingface")) |
| os.environ.setdefault("HF_HUB_DISABLE_XET", "1") |
| os.environ.setdefault("TOKENIZERS_PARALLELISM", "false") |
| os.environ.setdefault("GRADIO_ANALYTICS_ENABLED", "false") |
|
|
|
|
| def conversation(prompt: str, history: list[dict] | None = None) -> list[dict]: |
| if not prompt.strip(): |
| raise ValueError("Enter a message first.") |
| if len(prompt) > 24000: |
| raise ValueError("Please split this message into smaller parts.") |
| messages = [{"role": "system", "content": CONFIG["system_prompt"]}] |
| previous = history or [] |
| if len(previous) % 2: |
| raise ValueError("Conversation history is incomplete. Start a new chat.") |
| for index, item in enumerate(previous): |
| if item.get("role") != ("user" if index % 2 == 0 else "assistant"): |
| raise ValueError("Conversation history is invalid. Start a new chat.") |
| if not isinstance(item.get("content"), str): |
| raise ValueError("This chat accepts text messages only.") |
| messages.append({"role": item["role"], "content": item["content"]}) |
| messages.append({"role": "user", "content": prompt.strip()}) |
| return messages |
|
|
|
|
| def gpu_available() -> bool: |
| import torch |
| return bool(torch.cuda.is_available()) |
|
|
|
|
| def load_engine(adapter: str | None = None): |
| import torch |
| from transformers import AutoModelForCausalLM, AutoTokenizer |
|
|
| if not gpu_available(): |
| raise RuntimeError("The cloud model is not connected. This computer shows the interface preview only.") |
| tokenizer = AutoTokenizer.from_pretrained( |
| CONFIG["base_model"], revision=CONFIG["base_revision"], trust_remote_code=False) |
| model = AutoModelForCausalLM.from_pretrained( |
| CONFIG["base_model"], revision=CONFIG["base_revision"], trust_remote_code=False, |
| dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16, |
| device_map={"": "cuda:0"}, attn_implementation="sdpa") |
| if adapter: |
| from peft import PeftConfig, PeftModel |
| adapter_config = PeftConfig.from_pretrained(adapter) |
| if (adapter_config.base_model_name_or_path != CONFIG["base_model"] |
| or adapter_config.revision != CONFIG["base_revision"]): |
| raise ValueError("Adapter does not match the pinned chat-and-code model.") |
| model = PeftModel.from_pretrained(model, adapter) |
| return tokenizer, model.eval() |
|
|
|
|
| def answer(tokenizer, model, prompt: str, history: list[dict] | None = None, |
| max_new_tokens: int = CONFIG["max_new_tokens"]) -> str: |
| if not 1 <= max_new_tokens <= 2048: |
| raise ValueError("Response length must be between 1 and 2,048 tokens.") |
| messages = conversation(prompt, history) |
| while True: |
| inputs = tokenizer.apply_chat_template( |
| messages, tokenize=True, add_generation_prompt=True, |
| return_dict=True, return_tensors="pt") |
| if inputs["input_ids"].shape[1] + max_new_tokens <= 4096: |
| break |
| if len(messages) <= 2: |
| raise ValueError("This message is too long. Please send a smaller code example.") |
| del messages[1:3] |
| return generate_messages(tokenizer, model, messages, max_new_tokens=max_new_tokens) |
|
|
|
|
| def generate_messages(tokenizer, model, messages, max_new_tokens=768, |
| context_limit=4096, do_sample=True): |
| import torch |
| if not 1 <= max_new_tokens <= 2048 or not 256 <= context_limit <= 8192: |
| raise ValueError("Invalid response or context length.") |
| inputs = tokenizer.apply_chat_template( |
| messages, tokenize=True, add_generation_prompt=True, |
| return_dict=True, return_tensors="pt") |
| if inputs["input_ids"].shape[1] + max_new_tokens > context_limit: |
| raise ValueError("Source packet is too long. Shorten the excerpts; no sources were silently removed.") |
| inputs = inputs.to(model.device) |
| sampling = {"temperature": 0.7, "top_p": 0.8, "top_k": 20} if do_sample else {} |
| with torch.inference_mode(): |
| output = model.generate( |
| **inputs, max_new_tokens=max_new_tokens, do_sample=do_sample, **sampling, |
| pad_token_id=tokenizer.pad_token_id, eos_token_id=tokenizer.eos_token_id) |
| reply = tokenizer.decode(output[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True).strip() |
| if not reply: |
| raise RuntimeError("The model returned an empty response. Please try again.") |
| return reply |
|
|
