@lzhzzzzwill/cofos 1.0.2 → 1.1.0

package/runtime/src/inference/student_adapter_inference.py

@@ -11,16 +11,130 @@ 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."
+        "Student inference requires transformers and torch. "
+        "Install the inference dependencies first."
     ) from exc
 
+from data_processing.utils import resolve_under_root
 from inference.query_router import QueryRouter
-from inference.run_kg_grounded_inference import KGRAGInference
+from retrieval.retrieve_evidence import EvidenceRetriever
+
+
+class StudentRAGContext:
+    """Lightweight RAG helper for student-model inference.
+
+    This intentionally avoids the legacy teacher/LLM inference path; it only
+    loads retrieval indexes and formats retrieved context for Willlzh/COFOS.
+    """
+
+    def __init__(self, config: Dict[str, Any]):
+        self.retriever = EvidenceRetriever(config)
+
+    def format_kg_facts(self, kg_facts: List[Dict[str, Any]], max_items: int = 10) -> str:
+        if not kg_facts:
+            return "No relevant knowledge graph facts found."
+
+        support_icon = {
+            "supported": "✓",
+            "partially_supported": "~",
+            "unsupported": "?",
+        }
+        formatted = []
+        for fact in kg_facts[:max_items]:
+            head = fact.get("head", "")
+            relation = fact.get("relation", "")
+            tail = fact.get("tail", "")
+            tail_category = fact.get("tail_category", "")
+
+            line = f"- {head} {relation} {tail}"
+            if tail_category:
+                line += f" ({tail_category})"
+
+            condition = fact.get("condition") or {}
+            if isinstance(condition, dict):
+                condition_text = "; ".join(
+                    str(value)
+                    for value in [
+                        condition.get("phase"),
+                        condition.get("oxidant_source"),
+                        condition.get("driving_force"),
+                    ]
+                    if value
+                )
+                if condition_text:
+                    line += f" | condition: {condition_text}"
+
+            evidence = fact.get("evidence") or {}
+            if isinstance(evidence, dict):
+                status = evidence.get("support_status") or ""
+                level = evidence.get("evidence_level") or ""
+                evidence_text = " ".join(filter(None, [status, level]))
+                icon = support_icon.get(status, "")
+                if icon:
+                    evidence_text = f"{icon} {evidence_text}"
+            else:
+                evidence_text = str(evidence)
+            if evidence_text:
+                line += f" | {evidence_text}"
+
+            formatted.append(line)
+
+        return "\n".join(formatted)
+
+    @staticmethod
+    def _truncate_at_sentence(text: str, max_chars: int = 500) -> str:
+        if len(text) <= max_chars:
+            return text
+        snippet = text[:max_chars]
+        for sep in (". ", ".\n", "? ", "?\n", "! ", "!\n", "。", "；"):
+            idx = snippet.rfind(sep)
+            if idx > max_chars * 0.5:
+                return snippet[:idx + len(sep.rstrip())]
+        last_space = snippet.rfind(" ")
+        if last_space > max_chars * 0.5:
+            return snippet[:last_space]
+        return snippet
+
+    def format_evidence(self, evidence_list: List[Dict[str, Any]], max_items: int = 10) -> str:
+        if not evidence_list:
+            return "No relevant evidence found."
+
+        level_badge = {
+            "high": "[HIGH confidence]",
+            "medium": "[MEDIUM confidence]",
+            "low": "[LOW confidence]",
+        }
+        support_badge = {
+            "supported": "verified",
+            "partially_supported": "partial",
+            "unsupported": "unverified",
+        }
+
+        formatted = []
+        for i, evidence in enumerate(evidence_list[:max_items], 1):
+            source_type = evidence.get("source_type", "unknown")
+            material = evidence.get("material_context", {}).get("material_type", "N/A")
+            dominant = [
+                item.get("species", "N/A")
+                for item in evidence.get("dominant_ros_or_product", [])
+            ]
+            evidence_level = evidence.get("evidence_level", "low")
+            support_status = evidence.get("support_status", "unsupported")
+
+            level_str = level_badge.get(evidence_level, evidence_level)
+            support_str = support_badge.get(support_status, support_status)
+            text = self._truncate_at_sentence(evidence.get("text", ""), max_chars=500)
+
+            formatted.append(
+                f"[{i}] {source_type} {level_str} ({support_str})\n"
+                f"    Material: {material}\n"
+                f"    Dominant ROS: {', '.join(dominant)}\n"
+                f"    {text}"
+            )
+        return "\n\n".join(formatted)
 
 
 class StudentAdapterInference:
@@ -48,7 +162,12 @@ class StudentAdapterInference:
         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
+        if adapter_path:
+            raise ValueError(
+                "The standalone fwdemo runtime only supports merged models; "
+                "publish or pass a merged Transformers checkpoint instead of a LoRA adapter."
+            )
+        self.adapter_path = None
         self.model_is_hf = self._is_hf_repo(model_path)
         self.use_rag = use_rag
         generation_config = config.get("generation", {})
@@ -64,6 +183,7 @@ class StudentAdapterInference:
         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.stream_refinement = bool(generation_config.get("stream_refinement", True))
         self.refinement_max_new_tokens = int(
             generation_config.get(
                 "refinement_max_new_tokens",
@@ -71,12 +191,12 @@ class StudentAdapterInference:
             )
         )
         self.prompt_templates = self._load_prompt_templates()
-        self.rag = KGRAGInference(config) if use_rag else None
+        self.rag = StudentRAGContext(config) if use_rag else None
         self.tokenizer = None
         self.model = None
 
     def _load_prompt_templates(self) -> Dict[str, str]:
-        prompt_path = Path(
+        prompt_path = resolve_under_root(
             self.config.get("prompts", {}).get(
                 "path", "config/prompt_templates.yaml"
             )
@@ -101,18 +221,21 @@ class StudentAdapterInference:
             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}"
-            )
+
+    @staticmethod
+    def _select_dtype() -> torch.dtype:
+        """Pick the fastest dtype the current device actually supports.
+
+        bfloat16 was hardcoded, which silently falls back to a very slow
+        emulated path on CPU and on GPUs without bf16 support.
+        """
+        if torch.cuda.is_available():
+            if torch.cuda.is_bf16_supported():
+                return torch.bfloat16
+            return torch.float16
+        if torch.backends.mps.is_available():
+            return torch.float16
+        return torch.float32
 
     def load_model(self) -> None:
         self.validate_paths()
@@ -135,19 +258,21 @@ class StudentAdapterInference:
             )
         if self.tokenizer.pad_token is None:
             self.tokenizer.pad_token = self.tokenizer.eos_token
+        # Decoder-only models must be left-padded for correct generation.
+        self.tokenizer.padding_side = "left"
 
+        dtype = self._select_dtype()
+        logging.info("Loading model weights with dtype=%s", dtype)
         model = AutoModelForCausalLM.from_pretrained(
             self.model_path_str,
             device_map="auto",
-            torch_dtype=torch.bfloat16,
+            torch_dtype=dtype,
+            low_cpu_mem_usage=True,
             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")
+        logging.info("Using merged model without a separate LoRA adapter")
 
+        model.generation_config.use_cache = True
         self.model = model
         self.model.eval()
 
@@ -290,7 +415,7 @@ class StudentAdapterInference:
     ) -> str:
         prompt = self._render_prompt(messages)
         inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
-        with torch.no_grad():
+        with torch.inference_mode():
             output = self.model.generate(
                 **inputs,
                 **self._generation_kwargs(max_new_tokens=max_new_tokens),
@@ -298,6 +423,104 @@ class StudentAdapterInference:
         generated = output[0][inputs["input_ids"].shape[-1] :]
         return self.tokenizer.decode(generated, skip_special_tokens=True)
 
+    def _stream_from_messages(
+        self,
+        messages: List[Dict[str, str]],
+        max_new_tokens: Optional[int] = None,
+    ) -> Generator[str, None, None]:
+        """Yield generated text incrementally, suppressing reasoning scaffold.
+
+        Raw token streaming would leak ``<think>`` blocks and leading
+        ``Assistant:`` / ``Final Answer:`` labels to the user, because
+        :meth:`clean_answer` can only run on a complete string.  This wraps the
+        stream in a small state machine that withholds text until it is known to
+        be part of the real answer.
+        """
+        prompt = self._render_prompt(messages)
+        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(max_new_tokens=max_new_tokens),
+        }
+
+        errors: List[BaseException] = []
+
+        def _run() -> None:
+            try:
+                with torch.inference_mode():
+                    self.model.generate(**gen_kwargs)
+            except BaseException as exc:  # surfaced to the caller below
+                errors.append(exc)
+                streamer.end()
+
+        thread = Thread(target=_run, daemon=True)
+        thread.start()
+        try:
+            yield from self._sanitize_stream(streamer)
+        finally:
+            thread.join()
+        if errors:
+            raise errors[0]
+
+    @staticmethod
+    def _sanitize_stream(chunks: Any) -> Generator[str, None, None]:
+        """Drop ``<think>`` blocks and leading labels from a text stream."""
+        buffer = ""
+        started = False
+        in_think = False
+
+        for chunk in chunks:
+            buffer += chunk
+
+            while True:
+                if in_think:
+                    end = buffer.find("</think>")
+                    if end == -1:
+                        # Keep only enough tail to match a split closing tag.
+                        buffer = buffer[-len("</think>") :]
+                        break
+                    buffer = buffer[end + len("</think>") :]
+                    in_think = False
+                    continue
+
+                start = buffer.find("<think>")
+                if start != -1:
+                    buffer = buffer[start + len("<think>") :]
+                    in_think = True
+                    continue
+                break
+
+            if in_think:
+                continue
+
+            if not started:
+                # Hold back until we are past any leading label, and keep a tail
+                # that could still turn out to be the start of a "<think>" tag.
+                stripped = re.sub(
+                    r"(?is)^\s*(?:assistant|model|模型|final\s+answer|answer)\s*[:：]\s*",
+                    "",
+                    buffer,
+                )
+                if not stripped.strip():
+                    continue
+                buffer = stripped.lstrip()
+                started = True
+
+            if len(buffer) > len("<think>"):
+                emit, buffer = buffer[: -len("<think>")], buffer[-len("<think>") :]
+                if emit:
+                    yield emit
+
+        if not in_think and buffer.strip():
+            yield buffer if started else buffer.lstrip()
+
     @staticmethod
     def _has_report_style(answer: str) -> bool:
         report_patterns = (
@@ -316,26 +539,18 @@ class StudentAdapterInference:
         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:
+        messages = self._refinement_messages(question, draft_answer)
+        if messages is None:
             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},
-        ]
+            messages = [
+                messages[0],
+                {
+                    "role": "user",
+                    "content": f"{messages[1]['content']}\n\n{extra_instruction.strip()}",
+                },
+            ]
         return self.clean_answer(
             self._generate_from_messages(
                 messages,
@@ -399,13 +614,44 @@ class StudentAdapterInference:
     def answer(self, question: str, top_k: int = 5) -> str:
         return self.answer_with_metadata(question, top_k=top_k)["answer"]
 
+    def _refinement_messages(
+        self,
+        question: str,
+        draft_answer: str,
+    ) -> Optional[List[Dict[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 None
+        return [
+            {"role": "system", "content": system_prompt.strip()},
+            {
+                "role": "user",
+                "content": user_prompt_template.format(
+                    question=question,
+                    draft_answer=draft_answer,
+                ).strip(),
+            },
+        ]
+
     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`."""
+        """Stream the answer for *question* incrementally.
+
+        With ``refine_answers`` on, the draft pass is generated silently and the
+        *refinement* pass is the one streamed, so the user sees the final text
+        appear token by token rather than waiting for two full generations.  Set
+        ``generation.stream_refinement: false`` to keep the stricter
+        non-streaming path, which can additionally retry a report-style answer.
+        """
         quick_response = QueryRouter.quick_response(question)
         if quick_response is not None:
             yield quick_response
@@ -414,7 +660,7 @@ class StudentAdapterInference:
         if self.model is None or self.tokenizer is None:
             self.load_model()
 
-        if self.refine_answers:
+        if self.refine_answers and not self.stream_refinement:
             yield self.answer_with_metadata(
                 question,
                 top_k=top_k,
@@ -423,24 +669,29 @@ class StudentAdapterInference:
             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)
+        messages = self._with_history(context["messages"], history)
 
-        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()
+        if not self.refine_answers:
+            yield from self._stream_from_messages(messages)
+            return
+
+        draft = self.clean_answer(self._generate_from_messages(messages))
+        refinement = self._refinement_messages(question, draft)
+        if refinement is None:
+            yield draft
+            return
 
-        for text in streamer:
-            yield text
+        streamed = ""
+        for chunk in self._stream_from_messages(
+            refinement,
+            max_new_tokens=self.refinement_max_new_tokens,
+        ):
+            streamed += chunk
+            yield chunk
+
+        if not streamed.strip():
+            # The refinement pass produced nothing usable; fall back to the draft.
+            yield draft
 
     @staticmethod
     def _strip_reasoning_scaffold(answer: str) -> str:

package/runtime/src/retrieval/__init__.py

@@ -1,10 +1,17 @@
 #!/usr/bin/env python3
-"""
-COF-ROS-Reasoner: 检索模块初始化
-"""
+"""COFOS retrieval package."""
 
-from .build_faiss_index import FAISSIndexBuilder
-from .build_bm25_index import BM25IndexBuilder
-from .retrieve_evidence import EvidenceRetriever
 
-__all__ = ['FAISSIndexBuilder', 'BM25IndexBuilder', 'EvidenceRetriever']
+def __getattr__(name):
+    if name == "BM25IndexBuilder":
+        from .build_bm25_index import BM25IndexBuilder
+
+        return BM25IndexBuilder
+    if name == "EvidenceRetriever":
+        from .retrieve_evidence import EvidenceRetriever
+
+        return EvidenceRetriever
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+__all__ = ["BM25IndexBuilder", "EvidenceRetriever"]