Spaces:
Running
Running
| import re | |
| import logging | |
| from typing import List | |
| logger = logging.getLogger("axiom.postprocessor") | |
| class PostProcessor: | |
| """ | |
| Cleans generated text and injects source citations. | |
| """ | |
| def __init__(self, groq_gen=None): | |
| self.groq_gen = groq_gen | |
| logger.info("Ready.") | |
| def _remove_repetition(self, text: str) -> str: | |
| """Removes HTML artifacts and duplicate sentences.""" | |
| # Strip any HTML tags that leaked through | |
| text = re.sub(r'<[^>]+>', ' ', text) | |
| text = re.sub(r'\s+', ' ', text).strip() | |
| sentences = text.split(". ") | |
| seen = set() | |
| cleaned = [] | |
| for sentence in sentences: | |
| normalized = sentence.strip().lower() | |
| if normalized and normalized not in seen: | |
| seen.add(normalized) | |
| cleaned.append(sentence.strip()) | |
| return ". ".join(cleaned) | |
| def _inject_citations(self, text: str, sources: List[str]) -> str: | |
| """Appends source list at the bottom of the generated text.""" | |
| if not sources: | |
| return text | |
| citation_block = "\n\nSources:\n" | |
| for i, source in enumerate(sources): | |
| citation_block += f" [{i+1}] {source}\n" | |
| return text + citation_block | |
| def _check_faithfulness(self, generated: str, context: str) -> dict: | |
| """ | |
| LLM-based faithfulness check. | |
| """ | |
| if not self.groq_gen: | |
| return { | |
| "faithfulness_score": 0.5, | |
| "supported_sentences": 0, | |
| "unsupported_sentences": 0 | |
| } | |
| # Truncate context to avoid token limits causing failures | |
| truncated_context = context[:3000] if len(context) > 3000 else context | |
| truncated_generated = generated[:1000] if len(generated) > 1000 else generated | |
| prompt = f"""You are a strict fact-checker. Evaluate if the Generated Text is faithful to the Context. | |
| Context: | |
| {truncated_context} | |
| Generated Text: | |
| {truncated_generated} | |
| Rules: | |
| - faithfulness_score: 1.0 means every claim is supported by context | |
| - faithfulness_score: 0.0 means major claims are not in context | |
| - Count each sentence as supported or unsupported | |
| Respond ONLY with valid JSON, no markdown, no explanation: | |
| {{"faithfulness_score": 0.95, "supported_sentences": 4, "unsupported_sentences": 0}}""" | |
| try: | |
| res = self.groq_gen.generate(query=prompt, context="") | |
| # More aggressive cleaning | |
| clean_res = res.strip() | |
| clean_res = re.sub(r'```json\s*', '', clean_res) | |
| clean_res = re.sub(r'```\s*', '', clean_res) | |
| clean_res = clean_res.strip() | |
| # Find JSON object in response even if there's extra text | |
| json_match = re.search(r'\{[^}]+\}', clean_res) | |
| if json_match: | |
| clean_res = json_match.group() | |
| import json | |
| data = json.loads(clean_res) | |
| score = float(data.get("faithfulness_score", 0.5)) | |
| # Clamp between 0 and 1 | |
| score = max(0.0, min(1.0, score)) | |
| return { | |
| "faithfulness_score": score, | |
| "supported_sentences": int(data.get("supported_sentences", 0)), | |
| "unsupported_sentences": int(data.get("unsupported_sentences", 0)) | |
| } | |
| except Exception as e: | |
| logger.error(f"Faithfulness check failed: {e}") | |
| # Return 0.5 instead of 0.0 so badge still shows | |
| return { | |
| "faithfulness_score": 0.5, | |
| "supported_sentences": 0, | |
| "unsupported_sentences": 0 | |
| } | |
| def process(self, generated: str, sources: List[str], context: str) -> dict: | |
| """ | |
| Full post-processing pipeline. | |
| Returns cleaned text + citations + faithfulness check. | |
| """ | |
| cleaned = self._remove_repetition(generated) | |
| with_citations = self._inject_citations(cleaned, sources) | |
| faithfulness = self._check_faithfulness(cleaned, context) | |
| return { | |
| "final_output": with_citations, | |
| "clean_output": cleaned, | |
| "sources": sources, | |
| "faithfulness": faithfulness | |
| } |