Spaces:
Runtime error
Runtime error
File size: 6,252 Bytes
4e8fb20 cc3d1d7 4e8fb20 9899af7 4e8fb20 cc3d1d7 4e8fb20 cc3d1d7 4e8fb20 | 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 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 | import asyncio
import re
import threading
from typing import Any
import yaml
from config import HF_LOCAL_FILES_ONLY, TOPIC_PATTERN_BASE_MODEL, TOPIC_PATTERN_MODEL, USE_TOPIC_PATTERN_MODEL
class TopicPatternAgent:
"""Fine-tuned topic/pattern + coaching-step classifier."""
def __init__(
self,
frameworks_path: str = "frameworks.yaml",
model_name: str = TOPIC_PATTERN_MODEL,
base_model_name: str = TOPIC_PATTERN_BASE_MODEL,
):
self.model_name = model_name
self.base_model_name = base_model_name
self.enabled = USE_TOPIC_PATTERN_MODEL
self._model: Any | None = None
self._tokenizer: Any | None = None
self._device: str = "cpu"
self._load_error: str = ""
self.last_error = ""
self._load_lock = threading.Lock()
with open(frameworks_path, "r", encoding="utf-8") as file:
self.frameworks: dict[str, dict[str, Any]] = yaml.safe_load(file)
@property
def is_loaded(self) -> bool:
return self._model is not None and self._tokenizer is not None
async def warmup(self) -> None:
if not self.enabled:
self.last_error = "Topic/pattern model is disabled."
return
await asyncio.to_thread(self._ensure_model_loaded_sync)
async def analyze(self, question: str) -> dict[str, Any]:
if not self.enabled or not question.strip():
self.last_error = "Topic/pattern model is disabled or question is empty."
return {}
try:
output = await asyncio.to_thread(self._generate, question.strip())
except Exception as exc:
self._load_error = str(exc)
self.last_error = str(exc)
return {}
parsed = self.parse_output(output)
framework = self._valid_framework(parsed.get("type", ""))
steps = parsed.get("steps") or []
if not framework or not steps:
self.last_error = f"Could not parse topic model output: {output[:300]}"
return {}
self.last_error = ""
return {
"type": framework,
"pattern": parsed.get("type", ""),
"steps": steps,
"confidence": 0.85,
"model": self.model_name,
}
def parse_output(self, text: str) -> dict[str, Any]:
clean = self._extract_assistant_text(text)
type_match = re.search(r"Type\s*:\s*(.+?)(?:\n|$)", clean, flags=re.IGNORECASE)
steps_match = re.search(r"Steps\s*:\s*(.+)", clean, flags=re.IGNORECASE | re.DOTALL)
raw_type = type_match.group(1).strip() if type_match else ""
raw_steps = steps_match.group(1).strip() if steps_match else ""
steps = self._split_steps(raw_steps)
return {"type": raw_type, "steps": steps}
def _generate(self, question: str) -> str:
self._ensure_model_loaded_sync()
prompt = f"<|im_start|>user\n{question}<|im_end|>\n<|im_start|>assistant\n"
inputs = self._tokenizer(prompt, return_tensors="pt")
inputs = {key: value.to(self._model.device) for key, value in inputs.items()}
import torch
with torch.no_grad():
outputs = self._model.generate(
**inputs,
max_new_tokens=80,
temperature=0.1,
do_sample=False,
pad_token_id=self._tokenizer.eos_token_id,
)
return self._tokenizer.decode(outputs[0], skip_special_tokens=False)
def _ensure_model_loaded_sync(self) -> None:
if self._model is not None and self._tokenizer is not None:
return
with self._load_lock:
if self._model is not None and self._tokenizer is not None:
return
import torch
from huggingface_hub import snapshot_download
from transformers import AutoModelForCausalLM, AutoTokenizer
adapter_path = snapshot_download(self.model_name, local_files_only=HF_LOCAL_FILES_ONLY)
tokenizer = AutoTokenizer.from_pretrained(
self.base_model_name,
trust_remote_code=False,
local_files_only=HF_LOCAL_FILES_ONLY,
)
if tokenizer.pad_token_id is None:
tokenizer.pad_token = tokenizer.eos_token
model_kwargs = {"trust_remote_code": False, "low_cpu_mem_usage": True}
if torch.backends.mps.is_available():
model_kwargs["torch_dtype"] = torch.float16
try:
from peft import PeftModel
base_model = AutoModelForCausalLM.from_pretrained(
self.base_model_name,
local_files_only=HF_LOCAL_FILES_ONLY,
**model_kwargs,
)
model = PeftModel.from_pretrained(base_model, adapter_path)
except Exception as exc:
self.last_error = f"PEFT load failed: {exc}"
raise
if torch.backends.mps.is_available():
model = model.to("mps")
model.eval()
self._tokenizer = tokenizer
self._model = model
def _extract_assistant_text(self, text: str) -> str:
if "<|im_start|>assistant" in text:
text = text.split("<|im_start|>assistant", 1)[-1]
if "<|im_end|>" in text:
text = text.split("<|im_end|>", 1)[0]
return text.strip()
def _split_steps(self, text: str) -> list[str]:
if not text:
return []
parts = re.split(r"\s*(?:→|->|,|\n|;)\s*", text)
steps = []
for part in parts:
clean = re.sub(r"^\s*(?:[-*]|\d+[.)])\s*", "", part).strip()
if clean:
steps.append(clean)
return steps[:6]
def _valid_framework(self, value: str) -> str:
aliases = {
"Behavioural": "Behavioral",
"Product Design": "Product Sense",
"Data Science": "Technical",
"AI Engineering": "Technical",
"Estimation": "Case",
}
normalized = aliases.get(value.strip(), value.strip())
return normalized if normalized in self.frameworks else ""
|