@lzhzzzzwill/cofos 1.0.1 → 1.0.2

package/runtime/src/inference/student_adapter_inference.py

@@ -0,0 +1,520 @@
+#!/usr/bin/env python3
+"""Local student-model inference with optional confidence-gated KG/BM25 RAG."""
+
+import logging
+import re
+from pathlib import Path
+from threading import Thread
+from typing import Any, Dict, Generator, List, Optional
+
+import torch
+import yaml
+
+try:
+    from peft import PeftModel
+    from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+except ImportError as exc:
+    raise ImportError(
+        "Student inference requires transformers, peft, and torch. "
+        "Install the training/inference dependencies first."
+    ) from exc
+
+from inference.query_router import QueryRouter
+from inference.run_kg_grounded_inference import KGRAGInference
+
+
+class StudentAdapterInference:
+    """Run a fine-tuned student model with optional grounded retrieval."""
+
+    REQUIRED_PROMPTS = (
+        "student_qa_system_prompt",
+        "student_qa_user_prompt",
+        "student_rag_user_prompt",
+    )
+
+    @staticmethod
+    def _is_hf_repo(path: str) -> bool:
+        """Return True if ``path`` looks like a Hugging Face repo id (org/repo)."""
+        return bool(re.match(r"^[A-Za-z0-9_\-\.]+/[A-Za-z0-9_\-\.]+$", path))
+
+    def __init__(
+        self,
+        config: Dict[str, Any],
+        model_path: str,
+        adapter_path: Optional[str],
+        use_rag: bool = True,
+        max_new_tokens: Optional[int] = None,
+    ):
+        self.config = config
+        self.model_path_str = model_path
+        self.model_path = Path(model_path)
+        self.adapter_path = Path(adapter_path) if adapter_path else None
+        self.model_is_hf = self._is_hf_repo(model_path)
+        self.use_rag = use_rag
+        generation_config = config.get("generation", {})
+        self.max_new_tokens = int(
+            max_new_tokens or generation_config.get("max_new_tokens", 1024)
+        )
+        self.do_sample = bool(generation_config.get("do_sample", False))
+        self.temperature = float(generation_config.get("temperature", 0.2))
+        self.top_p = float(generation_config.get("top_p", 0.9))
+        self.repetition_penalty = float(generation_config.get("repetition_penalty", 1.08))
+        self.num_beams = int(generation_config.get("num_beams", 1))
+        self.no_repeat_ngram_size = int(generation_config.get("no_repeat_ngram_size", 0))
+        self.min_new_tokens = int(generation_config.get("min_new_tokens", 0))
+        self.history_max_turns = int(generation_config.get("history_max_turns", 4))
+        self.refine_answers = bool(generation_config.get("refine_answers", True))
+        self.refinement_max_new_tokens = int(
+            generation_config.get(
+                "refinement_max_new_tokens",
+                min(max(self.max_new_tokens, 128), 512),
+            )
+        )
+        self.prompt_templates = self._load_prompt_templates()
+        self.rag = KGRAGInference(config) if use_rag else None
+        self.tokenizer = None
+        self.model = None
+
+    def _load_prompt_templates(self) -> Dict[str, str]:
+        prompt_path = Path(
+            self.config.get("prompts", {}).get(
+                "path", "config/prompt_templates.yaml"
+            )
+        )
+        if not prompt_path.exists():
+            raise FileNotFoundError(f"Prompt template file not found: {prompt_path}")
+
+        with open(prompt_path, "r", encoding="utf-8") as handle:
+            prompts = yaml.safe_load(handle) or {}
+
+        missing = [key for key in self.REQUIRED_PROMPTS if not prompts.get(key)]
+        if missing:
+            raise ValueError(
+                f"Missing required prompt templates in {prompt_path}: {', '.join(missing)}"
+            )
+        return prompts
+
+    def validate_paths(self) -> None:
+        if self.model_is_hf:
+            # Hugging Face repo id — skip local path checks; transformers will
+            # download (or use the local HF cache) automatically.
+            return
+        if not self.model_path.exists():
+            raise FileNotFoundError(f"Base model path not found: {self.model_path}")
+        if self.adapter_path is None:
+            return
+        if not self.adapter_path.exists():
+            raise FileNotFoundError(f"Adapter path not found: {self.adapter_path}")
+        if not (self.adapter_path / "adapter_config.json").exists():
+            raise FileNotFoundError(
+                f"LoRA adapter_config.json not found in: {self.adapter_path}"
+            )
+        if not (self.adapter_path / "adapter_model.safetensors").exists():
+            raise FileNotFoundError(
+                f"LoRA adapter_model.safetensors not found in: {self.adapter_path}"
+            )
+
+    def load_model(self) -> None:
+        self.validate_paths()
+        logging.info("Loading base model from %s", self.model_path_str)
+        try:
+            self.tokenizer = AutoTokenizer.from_pretrained(
+                self.model_path_str, trust_remote_code=True
+            )
+        except (AttributeError, ValueError) as exc:
+            logging.warning(
+                "Default tokenizer loading failed (%s); retrying with Qwen2 fast tokenizer compatibility",
+                exc,
+            )
+            self.tokenizer = AutoTokenizer.from_pretrained(
+                self.model_path_str,
+                trust_remote_code=True,
+                use_fast=True,
+                tokenizer_type="qwen2",
+                extra_special_tokens={},
+            )
+        if self.tokenizer.pad_token is None:
+            self.tokenizer.pad_token = self.tokenizer.eos_token
+
+        model = AutoModelForCausalLM.from_pretrained(
+            self.model_path_str,
+            device_map="auto",
+            torch_dtype=torch.bfloat16,
+            trust_remote_code=True,
+        )
+        if self.adapter_path is not None:
+            logging.info("Loading LoRA adapter from %s", self.adapter_path)
+            model = PeftModel.from_pretrained(model, str(self.adapter_path))
+        else:
+            logging.info("Using merged model without a separate LoRA adapter")
+
+        self.model = model
+        self.model.eval()
+
+    def _generation_kwargs(self, max_new_tokens: Optional[int] = None) -> Dict[str, Any]:
+        kwargs = {
+            "max_new_tokens": int(max_new_tokens or self.max_new_tokens),
+            "do_sample": self.do_sample,
+            "repetition_penalty": self.repetition_penalty,
+            "eos_token_id": self.tokenizer.eos_token_id,
+            "pad_token_id": self.tokenizer.pad_token_id,
+        }
+        if self.do_sample:
+            kwargs["temperature"] = self.temperature
+            kwargs["top_p"] = self.top_p
+        if self.num_beams > 1:
+            kwargs["num_beams"] = self.num_beams
+            kwargs["early_stopping"] = True
+        if self.no_repeat_ngram_size > 0:
+            kwargs["no_repeat_ngram_size"] = self.no_repeat_ngram_size
+        if self.min_new_tokens > 0:
+            kwargs["min_new_tokens"] = self.min_new_tokens
+        return kwargs
+
+    def _qa_messages(
+        self,
+        question: str,
+        no_high_confidence_evidence: bool = False,
+    ) -> List[Dict[str, str]]:
+        template_key = (
+            "student_qa_no_evidence_user_prompt"
+            if no_high_confidence_evidence
+            and self.prompt_templates.get("student_qa_no_evidence_user_prompt")
+            else "student_qa_user_prompt"
+        )
+        user_prompt = self.prompt_templates[template_key].format(question=question)
+        return [
+            {
+                "role": "system",
+                "content": self.prompt_templates["student_qa_system_prompt"].strip(),
+            },
+            {"role": "user", "content": user_prompt.strip()},
+        ]
+
+    def _rag_context(self, question: str, top_k: int = 5) -> Dict[str, Any]:
+        if self.rag is None:
+            return {
+                "messages": self._qa_messages(question),
+                "mode": "qa",
+                "kg_facts": [],
+                "evidence": [],
+            }
+
+        retriever = self.rag.retriever
+        if hasattr(retriever, "is_domain_query") and not retriever.is_domain_query(question):
+            return {
+                "messages": self._qa_messages(question),
+                "mode": "qa",
+                "kg_facts": [],
+                "evidence": [],
+            }
+
+        try:
+            kg_facts = retriever.retrieve_kg_facts(
+                question,
+                top_k,
+                min_confidence=getattr(retriever, "min_kg_confidence", None),
+            )
+        except TypeError:
+            kg_facts = retriever.retrieve_kg_facts(question, top_k)
+        evidence = [
+            item for item in retriever.retrieve_with_sources(question, top_k)
+            if not hasattr(retriever, "is_confident_result")
+            or retriever.is_confident_result(item)
+        ]
+
+        if not kg_facts and not evidence:
+            logging.info("No high-confidence RAG context; answering in QA mode")
+            return {
+                "messages": self._qa_messages(
+                    question,
+                    no_high_confidence_evidence=True,
+                ),
+                "mode": "qa",
+                "kg_facts": [],
+                "evidence": [],
+            }
+
+        kg_facts_text = self.rag.format_kg_facts(kg_facts, max_items=top_k)
+        evidence_text = self.rag.format_evidence(evidence, max_items=top_k)
+        user_prompt = self.prompt_templates["student_rag_user_prompt"].format(
+            question=question,
+            kg_facts=kg_facts_text,
+            evidence=evidence_text,
+        )
+        return {
+            "messages": [
+                {
+                    "role": "system",
+                    "content": self.prompt_templates["student_qa_system_prompt"].strip(),
+                },
+                {"role": "user", "content": user_prompt.strip()},
+            ],
+            "mode": "rag",
+            "kg_facts": kg_facts,
+            "evidence": evidence,
+        }
+
+    def build_messages(self, question: str, top_k: int = 5) -> List[Dict[str, str]]:
+        return self._rag_context(question, top_k=top_k)["messages"]
+
+    def _with_history(
+        self,
+        messages: List[Dict[str, str]],
+        history: Optional[List[Dict[str, str]]] = None,
+    ) -> List[Dict[str, str]]:
+        if not history:
+            return messages
+        max_items = max(0, self.history_max_turns) * 2
+        recent_history = history[-max_items:] if max_items else []
+        if not recent_history:
+            return messages
+        return [messages[0], *recent_history, *messages[1:]]
+
+    def _render_prompt(self, messages: List[Dict[str, str]]) -> str:
+        if self.tokenizer is None:
+            raise RuntimeError("Tokenizer must be loaded before rendering a chat prompt")
+        return self.tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=True,
+        )
+
+    def build_prompt(self, question: str, top_k: int = 5) -> str:
+        return self._render_prompt(self.build_messages(question, top_k=top_k))
+
+    def _generate_from_messages(
+        self,
+        messages: List[Dict[str, str]],
+        max_new_tokens: Optional[int] = None,
+    ) -> str:
+        prompt = self._render_prompt(messages)
+        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
+        with torch.no_grad():
+            output = self.model.generate(
+                **inputs,
+                **self._generation_kwargs(max_new_tokens=max_new_tokens),
+            )
+        generated = output[0][inputs["input_ids"].shape[-1] :]
+        return self.tokenizer.decode(generated, skip_special_tokens=True)
+
+    @staticmethod
+    def _has_report_style(answer: str) -> bool:
+        report_patterns = (
+            r"(?i)\bbased on (?:the )?(?:provided|given|available|supplied)",
+            r"(?i)\bprovided (?:information|text|context|excerpt|data)",
+            r"(?i)\bgiven (?:the )?(?:information|text|context|excerpt|data)",
+            r"(?i)\b(?:supported|confirmed|backed) by (?:high )?(?:efficiency|performance|observations?|data|evidence|metrics)",
+            r"(?i)\bexplicit (?:mechanistic pathways?|detection evidence|data) (?:are|is) not",
+            r"(?i)\bprovided excerpt\b",
+        )
+        return any(re.search(pattern, answer) for pattern in report_patterns)
+
+    def _render_final_answer(
+        self,
+        question: str,
+        draft_answer: str,
+        extra_instruction: str = "",
+    ) -> str:
+        system_prompt = self.prompt_templates.get(
+            "student_answer_refinement_system_prompt"
+        )
+        user_prompt_template = self.prompt_templates.get(
+            "student_answer_refinement_user_prompt"
+        )
+        if not system_prompt or not user_prompt_template:
+            return draft_answer
+
+        user_prompt = user_prompt_template.format(
+            question=question,
+            draft_answer=draft_answer,
+        ).strip()
+        if extra_instruction:
+            user_prompt = f"{user_prompt}\n\n{extra_instruction.strip()}"
+
+        messages = [
+            {"role": "system", "content": system_prompt.strip()},
+            {"role": "user", "content": user_prompt},
+        ]
+        return self.clean_answer(
+            self._generate_from_messages(
+                messages,
+                max_new_tokens=self.refinement_max_new_tokens,
+            )
+        )
+
+    def _refine_answer(self, question: str, draft_answer: str) -> str:
+        if not self.refine_answers:
+            return draft_answer
+
+        refined = self._render_final_answer(question, draft_answer)
+        if not refined or refined == "No answer was generated.":
+            return draft_answer
+
+        if self._has_report_style(refined):
+            retry = self._render_final_answer(
+                question,
+                refined,
+                extra_instruction=(
+                    "The previous answer still sounded like a report about source material. "
+                    "Rewrite it once more as a direct answer to the user. Keep the same facts, "
+                    "but remove source-attribution language, evidence-status language, and comments "
+                    "about what is or is not present in excerpts/data unless the user asked for evidence."
+                ),
+            )
+            if retry and retry != "No answer was generated.":
+                return retry
+
+        return refined
+
+    def answer_with_metadata(
+        self,
+        question: str,
+        top_k: int = 5,
+        history: Optional[List[Dict[str, str]]] = None,
+    ) -> Dict[str, Any]:
+        quick_response = QueryRouter.quick_response(question)
+        if quick_response is not None:
+            return {
+                "answer": quick_response,
+                "mode": "qa",
+                "kg_facts_used": 0,
+                "evidence_used": 0,
+            }
+
+        if self.model is None or self.tokenizer is None:
+            self.load_model()
+
+        context = self._rag_context(question, top_k=top_k)
+        messages = self._with_history(context["messages"], history)
+        answer = self.clean_answer(self._generate_from_messages(messages))
+        answer = self._refine_answer(question, answer)
+        return {
+            "answer": answer,
+            "mode": context["mode"],
+            "kg_facts_used": len(context["kg_facts"]),
+            "evidence_used": len(context["evidence"]),
+        }
+
+    def answer(self, question: str, top_k: int = 5) -> str:
+        return self.answer_with_metadata(question, top_k=top_k)["answer"]
+
+    def answer_stream(
+        self,
+        question: str,
+        top_k: int = 5,
+        history: Optional[List[Dict[str, str]]] = None,
+    ) -> Generator[str, None, None]:
+        """Stream tokens for *question* via :class:`TextIteratorStreamer`."""
+        quick_response = QueryRouter.quick_response(question)
+        if quick_response is not None:
+            yield quick_response
+            return
+
+        if self.model is None or self.tokenizer is None:
+            self.load_model()
+
+        if self.refine_answers:
+            yield self.answer_with_metadata(
+                question,
+                top_k=top_k,
+                history=history,
+            )["answer"]
+            return
+
+        context = self._rag_context(question, top_k=top_k)
+        prompt = self._render_prompt(self._with_history(context["messages"], history))
+        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
+
+        streamer = TextIteratorStreamer(
+            self.tokenizer,
+            skip_prompt=True,
+            skip_special_tokens=True,
+        )
+        gen_kwargs = {
+            **inputs,
+            "streamer": streamer,
+            **self._generation_kwargs(),
+        }
+        thread = Thread(target=self.model.generate, kwargs=gen_kwargs)
+        thread.start()
+
+        for text in streamer:
+            yield text
+
+    @staticmethod
+    def _strip_reasoning_scaffold(answer: str) -> str:
+        final_match = re.search(
+            r"(?ims)^\s*(?:final\s+(?:answer|response)|answer)\s*[:：]\s*(.+)$",
+            answer,
+        )
+        if final_match:
+            return final_match.group(1).strip()
+
+        scaffold_heading = re.compile(
+            r"(?im)^\s*(?:thinking\s+process|analysis|plan|draft\s+the\s+final\s+answer|"
+            r"refine\s+for\s+natural\s+flow|final\s+polish|check\s+against\s+constraints)\s*[:：]?\s*$"
+        )
+        if not scaffold_heading.search(answer):
+            return answer
+
+        paragraphs = [part.strip() for part in re.split(r"\n\s*\n", answer) if part.strip()]
+        content_paragraphs = []
+        meta_line = re.compile(
+            r"(?i)^\s*(?:\d+[.)]\s*)?(?:analy[sz]e|evaluate|draft|refine|final\s+polish|"
+            r"check|constraints?|task|input|output|preserve|return\s+only|start\s+directly)\b"
+        )
+        bullet_meta = re.compile(r"^\s*[*-]\s*(?:Input|Task|Constraints?|Material|Condition|Primary|Nature|Evidence|Mechanism)\s*:", re.I)
+        for paragraph in paragraphs:
+            lines = [line.strip() for line in paragraph.splitlines() if line.strip()]
+            if not lines:
+                continue
+            if scaffold_heading.match(lines[0]):
+                if len(lines) > 1 and re.match(r"(?i)^\s*final\s+polish\s*[:：]?\s*$", lines[0]):
+                    content_paragraphs.append("\n".join(lines[1:]).strip())
+                continue
+            if meta_line.match(lines[0]) or bullet_meta.match(lines[0]):
+                continue
+            content_paragraphs.append(paragraph)
+
+        return content_paragraphs[-1] if content_paragraphs else answer
+
+    @staticmethod
+    def clean_answer(answer: str) -> str:
+        answer = answer.strip()
+        if not answer:
+            return "No answer was generated."
+        answer = re.sub(r"(?is)<think>.*?</think>", "", answer).strip()
+        answer = re.sub(r"(?is)^.*?</think>", "", answer).strip()
+        answer = StudentAdapterInference._strip_reasoning_scaffold(answer)
+        if not answer:
+            return "No answer was generated."
+        answer = re.sub(r"(?im)^\s*(assistant|model|模型)\s*[:：]\s*", "", answer)
+        answer = re.sub(
+            r"(?is)^\s*sure[,，]?\s*(here(?:'s| is)\s+)?(?:the\s+)?answer\s*[:：]?\s*",
+            "",
+            answer,
+        )
+        return re.sub(r"\n{3,}", "\n\n", answer).strip()
+
+    def chat(self, top_k: int = 5) -> None:
+        print("进入 COF-ROS 学生模型问答模式，输入 exit 退出。")
+        history: List[Dict[str, str]] = []
+        while True:
+            question = input("\n你: ").strip()
+            if question.lower() in {"exit", "quit", "q"}:
+                print("退出")
+                break
+            if question.lower() in {"clear", "/clear"}:
+                history.clear()
+                print("已清空对话历史")
+                continue
+            if not question:
+                continue
+            print("\n模型:")
+            result = self.answer_with_metadata(question, top_k=top_k, history=history)
+            print(result["answer"])
+            history.extend([
+                {"role": "user", "content": question},
+                {"role": "assistant", "content": result["answer"]},
+            ])

package/runtime/src/retrieval/__init__.py

@@ -0,0 +1,10 @@
+#!/usr/bin/env python3
+"""
+COF-ROS-Reasoner: 检索模块初始化
+"""
+
+from .build_faiss_index import FAISSIndexBuilder
+from .build_bm25_index import BM25IndexBuilder
+from .retrieve_evidence import EvidenceRetriever
+
+__all__ = ['FAISSIndexBuilder', 'BM25IndexBuilder', 'EvidenceRetriever']