File size: 4,742 Bytes
5ba1bb4 18aeae0 5ba1bb4 18aeae0 5ba1bb4 18aeae0 5ba1bb4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 | """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
|