Upload 11 files

src/alignment.py

@@ -0,0 +1,256 @@
+"""
+Adapted from Jonathan Dursi
+https://github.com/ljdursi/poapy
+"""
+
+import numpy
+
+
+class SeqGraphAlignment(object):
+    __matchscore = 1
+    __mismatchscore = -2
+    __gap = -1
+
+    def __init__(
+        self,
+        sequence,
+        graph,
+        fastMethod=True,
+        globalAlign=False,
+        matchscore=__matchscore,
+        mismatchscore=__mismatchscore,
+        gapscore=__gap,
+        *args,
+        **kwargs,
+    ):
+        self._mismatchscore = mismatchscore
+        self._matchscore = matchscore
+        self._gap = gapscore
+        self.sequence = sequence
+        self.graph = graph
+        self.stringidxs = None
+        self.nodeidxs = None
+        self.globalAlign = globalAlign
+        if fastMethod:
+            matches = self.alignStringToGraphFast(*args, **kwargs)
+        else:
+            matches = self.alignStringToGraphSimple(*args, **kwargs)
+        self.stringidxs, self.nodeidxs = matches
+
+    def alignmentStrings(self):
+        return "".join(
+            self.sequence[i] if i is not None else "-" for i in self.stringidxs
+        ), "".join(self.graph.nodedict[j].text if j is not None else "-" for j in self.nodeidxs)
+
+    def matchscore(self, c1, c2):
+        if c1 == c2:
+            return self._matchscore
+        else:
+            return self._mismatchscore
+
+    def matchscoreVec(self, c, v):
+        return numpy.where(v == c, self._matchscore, self._mismatchscore)
+
+    def alignStringToGraphSimple(self):
+        """Align string to graph, following same approach as smith waterman
+        example"""
+        if type(self.sequence) is not str:
+            raise TypeError("Invalid Type")
+
+        nodeIDtoIndex, nodeIndexToID, scores, backStrIdx, backGrphIdx = (
+            self.initializeDynamicProgrammingData()
+        )
+
+        # Dynamic Programming
+        ni = self.graph.nodeiterator()
+        for i, node in enumerate(ni()):
+            pbase = node.text
+
+            for j, sbase in enumerate(self.sequence):
+                # add all candidates to a list, pick the best
+                candidates = [(scores[i + 1, j] + self._gap, i + 1, j, "INS")]
+                for predIndex in self.prevIndices(node, nodeIDtoIndex):
+                    candidates += [
+                        (scores[predIndex + 1, j + 1] + self._gap, predIndex + 1, j + 1, "DEL")
+                    ]
+                    candidates += [
+                        (
+                            scores[predIndex + 1, j] + self.matchscore(sbase, pbase),
+                            predIndex + 1,
+                            j,
+                            "MATCH",
+                        )
+                    ]
+
+                (
+                    scores[i + 1, j + 1],
+                    backGrphIdx[i + 1, j + 1],
+                    backStrIdx[i + 1, j + 1],
+                    movetype,
+                ) = max(candidates)
+
+                if not self.globalAlign and scores[i + 1, j + 1] < 0:
+                    scores[i + 1, j + 1] = 0.0
+                    backGrphIdx[i + 1, j + 1] = -1
+                    backStrIdx[i + 1, j + 1] = -1
+
+        return self.backtrack(scores, backStrIdx, backGrphIdx, nodeIndexToID)
+
+    def alignStringToGraphFast(self):
+        """Align string to graph - using numpy to vectorize across the string
+        at each iteration."""
+        if type(self.sequence) is not str:
+            raise TypeError("Invalid Type")
+
+        l2 = len(self.sequence)
+        seqvec = numpy.array(list(self.sequence))
+
+        nodeIDtoIndex, nodeIndexToID, scores, backStrIdx, backGrphIdx = (
+            self.initializeDynamicProgrammingData()
+        )
+        inserted = numpy.zeros((l2), dtype=bool)
+
+        # having the inner loop as a function improves performance
+        # can use Cython, etc on this for significant further improvements
+        # can't vectorize this since there's a loop-carried dependency
+        #  along the string
+        def insertions(i, l2, scores, inserted):
+            inserted[:] = False
+            for j in range(l2):
+                insscore = scores[i + 1, j] + self._gap
+                if insscore >= scores[i + 1, j + 1]:
+                    scores[i + 1, j + 1] = insscore
+                    inserted[j] = True
+
+        # Dynamic Programming
+        ni = self.graph.nodeiterator()
+        for i, node in enumerate(ni()):
+            gbase = node.text
+            predecessors = self.prevIndices(node, nodeIDtoIndex)
+
+            # calculate all best deletions, matches in one go over all
+            # predecessors.
+
+            # First calculate for the first predecessor, over all string posns:
+            deletescore = scores[predecessors[0] + 1, 1:] + self._gap
+            bestdelete = numpy.zeros((l2), dtype=numpy.int32) + predecessors[0] + 1
+
+            matchpoints = self.matchscoreVec(gbase, seqvec)
+            matchscore = scores[predecessors[0] + 1, 0:-1] + matchpoints
+            bestmatch = numpy.zeros((l2), dtype=numpy.int32) + predecessors[0] + 1
+
+            # then, the remaining
+            for predecessor in predecessors[1:]:
+                newdeletescore = scores[predecessor + 1, 1:] + self._gap
+                bestdelete = numpy.where(newdeletescore > deletescore, predecessor + 1, bestdelete)
+                deletescore = numpy.maximum(newdeletescore, deletescore)
+
+                gbase = self.graph.nodeIdxToBase(predecessor)
+                matchpoints = self.matchscoreVec(gbase, seqvec)
+                newmatchscore = scores[predecessor + 1, 0:-1] + matchpoints
+                bestmatch = numpy.where(newmatchscore > matchscore, predecessor + 1, bestmatch)
+                matchscore = numpy.maximum(newmatchscore, matchscore)
+
+            # choose best options available of match, delete
+            deleted = deletescore >= matchscore
+            backGrphIdx[i + 1, 1:] = numpy.where(deleted, bestdelete, bestmatch)
+            backStrIdx[i + 1, 1:] = numpy.where(
+                deleted, numpy.arange(1, l2 + 1), numpy.arange(0, l2)
+            )
+            scores[i + 1, 1:] = numpy.where(deleted, deletescore, matchscore)
+
+            # insertions: updated in place, don't depend on predecessors
+            insertions(i, l2, scores, inserted)
+            backGrphIdx[i + 1, 1:] = numpy.where(inserted, i + 1, backGrphIdx[i + 1, 1:])
+            backStrIdx[i + 1, 1:] = numpy.where(inserted, numpy.arange(l2), backStrIdx[i + 1, 1:])
+
+            # if we're doing local alignment, don't let bad global alignment
+            # drag us negative
+            if not self.globalAlign:
+                backGrphIdx[i + 1, :] = numpy.where(scores[i + 1, :] > 0, backGrphIdx[i + 1, :], -1)
+                backStrIdx[i + 1, :] = numpy.where(scores[i + 1, :] > 0, backStrIdx[i + 1, :], -1)
+                scores[i + 1, :] = numpy.maximum(scores[i + 1, :], 0)
+
+        return self.backtrack(scores, backStrIdx, backGrphIdx, nodeIndexToID)
+
+    def prevIndices(self, node, nodeIDtoIndex):
+        """Return a list of the previous dynamic programming table indices
+        corresponding to predecessors of the current node."""
+        prev = [nodeIDtoIndex[predID] for predID in list(node.inEdges.keys())]
+        # if no predecessors, point to just before the graph
+        if not prev:
+            prev = [-1]
+        return prev
+
+    def initializeDynamicProgrammingData(self):
+        """Initalize the dynamic programming tables:
+        - set up scores array
+        - set up backtracking array
+        - create index to Node ID table and vice versa"""
+        l1 = self.graph.nNodes
+        l2 = len(self.sequence)
+
+        nodeIDtoIndex = {}
+        nodeIndexToID = {-1: None}
+        # generate a dict of (nodeID) -> (index into nodelist (and thus matrix))
+        ni = self.graph.nodeiterator()
+        for index, node in enumerate(ni()):
+            nodeIDtoIndex[node.ID] = index
+            nodeIndexToID[index] = node.ID
+
+        # Dynamic Programming data structures; scores matrix and backtracking
+        # matrix
+        scores = numpy.zeros((l1 + 1, l2 + 1), dtype=numpy.int32)
+
+        # initialize insertion score
+        # if global align, penalty for starting at head != 0
+        if self.globalAlign:
+            scores[0, :] = numpy.arange(l2 + 1) * self._gap
+
+            ni = self.graph.nodeiterator()
+            for index, node in enumerate(ni()):
+                prevIdxs = self.prevIndices(node, nodeIDtoIndex)
+                best = scores[prevIdxs[0] + 1, 0]
+                for prevIdx in prevIdxs:
+                    best = max(best, scores[prevIdx + 1, 0])
+                scores[index + 1, 0] = best + self._gap
+
+        # backtracking matrices
+        backStrIdx = numpy.zeros((l1 + 1, l2 + 1), dtype=numpy.int32)
+        backGrphIdx = numpy.zeros((l1 + 1, l2 + 1), dtype=numpy.int32)
+
+        return nodeIDtoIndex, nodeIndexToID, scores, backStrIdx, backGrphIdx
+
+    def backtrack(self, scores, backStrIdx, backGrphIdx, nodeIndexToID):
+        """Backtrack through the scores and backtrack arrays.
+        Return a list of sequence indices and node IDs (not indices, which
+        depend on ordering)."""
+        besti, bestj = scores.shape
+        besti -= 1
+        bestj -= 1
+        if not self.globalAlign:
+            besti, bestj = numpy.argwhere(scores == numpy.amax(scores))[-1]
+        else:
+            ni = self.graph.nodeiterator()
+            # still have to find best final index to start from
+            terminalIndices = [index for (index, node) in enumerate(ni()) if node.outDegree == 0]
+            print(terminalIndices)
+            besti = terminalIndices[0] + 1
+            bestscore = scores[besti, bestj]
+            for i in terminalIndices[1:]:
+                score = scores[i + 1, bestj]
+                if score > bestscore:
+                    bestscore, besti = score, i + 1
+
+        matches = []
+        strindexes = []
+        while (self.globalAlign or scores[besti, bestj] > 0) and (besti != 0 or bestj != 0):
+            nexti, nextj = backGrphIdx[besti, bestj], backStrIdx[besti, bestj]
+            curstridx, curnodeidx = bestj - 1, nodeIndexToID[besti - 1]
+
+            strindexes.insert(0, curstridx if nextj != bestj else None)
+            matches.insert(0, curnodeidx if nexti != besti else None)
+
+            besti, bestj = nexti, nextj
+
+        return strindexes, matches

src/generation_methods.py

@@ -0,0 +1,299 @@
+from typing import Optional
+
+from src.generation_utils import (
+    extract_alternative_paths,
+    extract_context,
+    extract_equivalent_classes,
+    self_complete,
+    verify_correctness_pairwise,
+)
+from src.global_edit_utils import clean_up_text
+from src.text_poa_graph import TextPOAGraph
+
+"""
+Decodes from a TextPOAGraph object to a string by sequentially selecting nodes based on the selection threshold.
+Only the primary variation of selected variable nodes are selected.
+Text is edited using the global_edit_function (e.g. to clean up text by removing incoherencies, disfluencies, and redundancies).
+
+Args:
+    text_poa_graph: The TextPOAGraph object to decode.
+    selection_threshold: The threshold for selecting nodes.
+    model: The model to use for decoding.
+
+Returns:
+    A string of the decoded text.
+"""
+
+
+def decode_consensus(
+    text_poa_graph: TextPOAGraph,
+    selection_threshold: Optional[float] = 0.5,
+    task: str = "bio",
+    verbose: bool = False,
+    **kwargs,
+) -> str:
+    if text_poa_graph.failed:
+        return "Abstain"
+
+    text_poa_graph.toposort()
+
+    consensus_node_ids = text_poa_graph.consensus_node_ids
+
+    selected_node_ids = []
+
+    for node_id in consensus_node_ids:
+        if node_id == text_poa_graph.start_id or node_id == text_poa_graph.end_id:
+            continue
+
+        selected_node_ids.append(node_id)
+
+        for neighbor_id in text_poa_graph.nodedict[node_id].outEdges:
+            if neighbor_id in consensus_node_ids:
+                continue
+
+            if (
+                len(text_poa_graph.nodedict[neighbor_id].labels) / text_poa_graph.num_sequences
+                >= selection_threshold
+            ):
+                selected_node_ids.append(neighbor_id)
+
+    texts = []
+    for node_id in selected_node_ids:
+        if not text_poa_graph.nodedict[node_id].variations:
+            texts.append(text_poa_graph.nodedict[node_id].text)
+        else:
+            all_texts = [v for v in text_poa_graph.nodedict[node_id].variations.values()]
+            all_texts.append(text_poa_graph.nodedict[node_id].text)
+            # select the variation that is longest
+            texts.append(max(all_texts, key=len))
+    text = " ".join(texts)
+    edited_text = clean_up_text(text=text, task=task, api="openai", **kwargs)
+    if verbose:
+        return text, edited_text
+    else:
+        return edited_text
+
+
+def decode_self_verified(
+    text_poa_graph: TextPOAGraph,
+    problem: str,
+    uncertainty_threshold: float = 0.6,
+    verification_api: str = "openai",
+    verification_model: str = "gpt-4o-mini",
+    grace_period: bool = True,
+):
+    high_uncertainty_nodes = []
+    for node_id in text_poa_graph.consensus_node_ids:
+        if node_id == text_poa_graph.start_id or node_id == text_poa_graph.end_id:
+            continue
+
+        outgoing_edges = text_poa_graph.nodedict[node_id].outEdges
+        branching_factor = len(outgoing_edges) / text_poa_graph.num_sequences
+
+        if branching_factor > uncertainty_threshold:
+            high_uncertainty_nodes.append(node_id)
+
+    selected_labels = list(text_poa_graph._seq_paths.keys())
+    masked_candidates = {}
+    uncertain_region = False
+    for label in selected_labels:
+        text = ""
+        for node_id in text_poa_graph._seq_paths[label]:
+            if uncertain_region:
+                text += f" *START_SEPARATOR*_{node_id} "
+            if node_id in high_uncertainty_nodes:
+                uncertain_region = True
+
+            if len(text_poa_graph.nodedict[node_id].variations) > 0:
+                text += text_poa_graph.nodedict[node_id].variations[label]
+                text += " "
+            else:
+                text += text_poa_graph.nodedict[node_id].text
+                text += " "
+
+            if uncertain_region and node_id not in high_uncertainty_nodes:
+                text += f" *END_SEPARATOR*_{node_id} "
+                uncertain_region = False
+        masked_candidates[label] = text
+
+    patch_start_node = None
+    uncertain_ids = []
+
+    # give a grace period for the first incorrect step
+    prev_step = {label: None for label in selected_labels}
+
+    for node_id in high_uncertainty_nodes:
+        uncertain_ids.append(node_id)
+        context_before = extract_context(text_poa_graph, node_id)
+        alternative_paths = extract_alternative_paths(text_poa_graph, node_id)
+        equivalent_classes = extract_equivalent_classes(text_poa_graph, node_id, selected_labels)
+        new_labels = selected_labels.copy()
+
+        # Only do self-verifaction for labels from different sematically equivalent branches
+        if len(equivalent_classes) <= 1:
+            continue
+        i = 0
+        while i < len(equivalent_classes):
+            if i + 1 < len(equivalent_classes):
+                label_a = equivalent_classes[i][0]
+                label_b = equivalent_classes[i + 1][0]
+                full_a = context_before[label_a] + alternative_paths[label_a]
+                full_b = context_before[label_b] + alternative_paths[label_b]
+
+                score = verify_correctness_pairwise(
+                    full_text_1=full_a,
+                    full_text_2=full_b,
+                    verification_model=verification_model,
+                    problem=problem,
+                    api=verification_api,
+                )
+                if float(score[0]) < 1.0:
+                    print(f"Label {label_a} is incorrect at node {node_id}")
+                    masked_candidates[label_a] = (
+                        masked_candidates[label_a]
+                        .replace(f" *START_SEPARATOR*_{node_id} ", "*START_POSSIBLE_ERROR*")
+                        .replace(f" *END_SEPARATOR*_{node_id} ", "*END_POSSIBLE_ERROR*")
+                    )
+                    if not prev_step[label_a]:
+                        prev_step[label_a] = True
+                    if prev_step[label_a] and grace_period or not grace_period:
+                        for label_i in equivalent_classes[i]:
+                            new_labels.remove(label_i)
+                            print(f"\nSequence {label_i} pruned at node {node_id} (pairwise)")
+                if float(score[0]) == 1.0:
+                    prev_step[label_a] = False
+                if float(score[1]) < 1.0:
+                    print(f"Label {label_b} is incorrect at node {node_id}")
+                    masked_candidates[label_b] = (
+                        masked_candidates[label_b]
+                        .replace(f" *START_SEPARATOR*_{node_id} ", "*START_POSSIBLE_ERROR*")
+                        .replace(f" *END_SEPARATOR*_{node_id} ", "*END_POSSIBLE_ERROR*")
+                    )
+                    if not prev_step[label_b]:
+                        prev_step[label_b] = True
+                    if prev_step[label_b] and grace_period or not grace_period:
+                        for label_i in equivalent_classes[i + 1]:
+                            new_labels.remove(label_i)
+                            print(f"\nSequence {label_i} pruned at node {node_id} (pairwise)")
+                if float(score[1]) == 1.0:
+                    prev_step[label_b] = False
+                i += 2
+            else:
+                break
+
+        if len(new_labels) == 0:
+            patch_start_node = node_id
+            break
+
+        selected_labels = new_labels.copy()
+
+    # These are the pruned approaches with masking
+    print(masked_candidates)
+    masked_approaches = "\n".join(
+        [
+            f"Approach {label}: {masked_candidates[label].replace('START_SEPARATOR', 'START_UNCERTAIN_REGION').replace('END_SEPARATOR', 'END_UNCERTAIN_REGION')}"
+            for label in selected_labels
+        ]
+    )
+    # These are all approaches with masking
+    all_approaches = "\n".join(
+        [f"Approach {label}: {masked_candidates[label]}" for label in masked_candidates.keys()]
+    )
+
+    default_prompt = f"""
+    Solve the following math problem with mathematical precision and clarity.
+
+    Problem: {problem}
+
+    Below are potential solution approaches with sections marked as uncertain (between *START_UNCERTAIN_REGION* and *END_UNCERTAIN_REGION*). 
+    These sections may contain conceptual or computational errors.
+
+    There are also sections marked as *START_POSSIBLE_ERROR* and *END_POSSIBLE_ERROR*.
+    A verification step indicated that these steps are highly likely to contain errors.
+
+    Potential Approaches:
+    {masked_approaches}
+
+    Your task:
+    1. Analyze all potential approaches critically, identifying their mathematical strengths and weaknesses
+       If the approaches contain different answers, think carefully about why they are different, and use this to identify potential errors.
+    2. Using the sections with special markers, identify potential errors.
+    3. Develop a rigorous, step-by-step solution based on sound mathematical principles
+    4. For uncertain regions:
+       - Verify each step using algebraic or numerical validation
+       - If correct, incorporate these steps with appropriate justification
+       - If incorrect, provide clear corrections with mathematical reasoning for your changes
+    5. Follow a comparative approach, using the differences between approaches to identify potential errors.
+    6. Do not blindly follow the approaches, but rather use them to identify potential errors.
+
+    Guidelines for your solution:
+    - Begin with a strategic overview of your chosen approach
+    - Present each mathematical step with clear notation and justification
+    - Pay special attention to areas that were previously marked uncertain
+
+    Conclude your solution with:
+    Therefore, the final answer is: $\\boxed{{answer}}$.
+
+    Solution:
+    """
+
+    patch_prompt = f"""
+    Solve the following mathematical problem with precision and clarity.
+
+    Problem: {problem}
+
+    You have been provided with several partial solution approaches that attempted to solve this problem. 
+    None of these approaches are correct, but may contain valuable insights.
+    Sections marked between *START_POSSIBLE_ERROR* and *END_POSSIBLE_ERROR* indicate steps where previous solutions showed uncertainty.
+    A verification step indicated that these steps are likely to contain errors.
+
+    INSTRUCTIONS:
+    1. Synthesize a correct solution using insights from the previous approaches
+    2. Pay special attention to fixing the problematic areas marked by separators
+    3. Develop your solution step-by-step, showing clear mathematical reasoning
+    4. Focus especially on mathematical correctness in areas where previous solutions diverged
+    5. Present your work in a logical, sequential manner suitable for an advanced reader
+
+    GUIDELINES FOR MATHEMATICAL RIGOR:
+    1. MAINTAIN MATHEMATICAL RIGOR
+    - Verify that all mathematical operations follow from established principles and definitions
+    - Ensure dimensional consistency throughout calculations
+    - Check that algebraic manipulations preserve equality and do not introduce errors
+   
+    2. CONSIDER ALTERNATIVE PERSPECTIVES
+    - Even when approaches reach the same conclusion, examine their reasoning independently
+    - Look for more elegant or insightful connections that may be missed across all approaches
+    - Consider whether fundamental mathematical principles suggest a different path
+   
+    3. CRITICAL VALIDATION
+    - Test conclusions using known mathematical properties and relationships
+    - When possible, verify results using alternative methods
+    - Be especially cautious when all approaches agree on a result but use similar reasoning
+   
+    4. USE PRECISION IN CORRECTIONS
+    - When correcting uncertain regions, specify exactly what was incorrect and why
+    - Provide clear mathematical justification for any changes
+    - Ensure corrections align with standard mathematical principles and notations
+
+    Previous Approaches (for reference only):
+{all_approaches}
+
+Your Solution:
+[Begin with a clear statement of your approach]
+[Provide detailed mathematical steps]
+[Ensure correct handling of complex mathematical operations]
+[Verify your work at key points, especially in previously problematic areas]
+
+Always conclude with:
+Therefore, the final answer is: $\\boxed{{answer}}$ 
+    """
+
+    if patch_start_node is not None or len(masked_candidates.keys()) == 1:
+        print("None correct, patching")
+        prompt = patch_prompt
+    else:
+        prompt = default_prompt
+
+    return self_complete(
+        verification_prompt=prompt, verification_model=verification_model, api=verification_api
+    ), masked_candidates

src/generation_utils.py

@@ -0,0 +1,190 @@
+import re
+
+from huggingface_hub import InferenceClient
+from openai import OpenAI
+from together import Together
+
+from src.text_poa_graph import TextPOAGraph
+
+
+def extract_context(text_poa_graph, node_id):
+    """Extract context up to and including the specified node_id."""
+    contexts = {}
+    for label, path in text_poa_graph._seq_paths.items():
+        idx = path.index(node_id)
+        context = path[: idx + 1]
+        contexts[label] = " ".join(
+            text_poa_graph.nodedict[nid].variations.get(label, text_poa_graph.nodedict[nid].text)
+            for nid in context
+        )
+    return contexts
+
+
+def extract_alternative_paths(text_poa_graph: TextPOAGraph, node_id):
+    """Extract all alternative paths from this uncertainty point to the next consensus node."""
+    alternative_paths = {}
+    for label, path in text_poa_graph._seq_paths.items():
+        idx = path.index(node_id)
+        next_cn = None
+        for i in range(idx + 1, len(path)):
+            if path[i] in text_poa_graph.consensus_node_ids:
+                next_cn = path[i]
+                break
+
+        if next_cn:
+            next_cn_idx = path.index(next_cn)
+            alternative_segment = path[idx + 1 : next_cn_idx + 1]
+        else:
+            alternative_segment = []
+
+        alternative_paths[label] = " ".join(
+            text_poa_graph.nodedict[nid].variations.get(label, text_poa_graph.nodedict[nid].text)
+            for nid in alternative_segment
+        )
+    return alternative_paths
+
+
+def is_same_branch(text_poa_graph: TextPOAGraph, node_id, lable_1, label_2):
+    """Check if the next vaiable nodes for two sequences are the same after node_id."""
+    path_1 = text_poa_graph._seq_paths[lable_1]
+    path_2 = text_poa_graph._seq_paths[label_2]
+    idx_1 = path_1.index(node_id)
+    idx_2 = path_2.index(node_id)
+    return path_1[idx_1 + 1] == path_2[idx_2 + 1]
+
+
+def extract_equivalent_classes(text_poa_graph: TextPOAGraph, node_id, selected_labels):
+    """Extract equivalent classes from the text POA graph."""
+    if not selected_labels:
+        return []
+
+    equivalent_classes = []
+    for label in selected_labels:
+        matched = False
+        for class_group in equivalent_classes:
+            if is_same_branch(text_poa_graph, node_id, class_group[0], label):
+                class_group.append(label)
+                matched = True
+                break
+        if not matched:
+            equivalent_classes.append([label])
+    return equivalent_classes
+
+
+def verify_correctness_pairwise(
+    full_text_1: str, full_text_2: str, verification_model: str, problem: str, api: str = "openai"
+):
+    """Pairwise verification of two partial solution paths."""
+    if api == "openai":
+        client = OpenAI()
+    elif api == "hf":
+        client = InferenceClient()
+    elif api == "together":
+        client = Together()
+    else:
+        raise ValueError(f"Invalid API: {api}")
+
+    prompt = f"""
+     You will be given a problem and 2 partial solutions. 
+    Your task is to use comparison as an EFFICIENCY TOOL to quickly identify potential errors.
+    You will be given guidelines to follow, and you will be penalized if you do not follow them.
+
+    Problem: {problem}
+
+    Partial Solution 1: {full_text_1}
+    Partial Solution 2: {full_text_2}
+
+    CRITICAL GUIDELINES:
+    - DO NOT penalize a solution for being incomplete or having missing steps
+    - DO NOT make a comparison of which solution is better
+    - DO NOT consider steps incorrect just because they differ between solutions
+    - DO NOT prematurely evaluate based on final answers or future steps
+    - DO NOT expect both solutions to be at the same stage of completion
+    - DO NOT consider a step incorrect just because it lacks sufficient detail or justification
+
+    KEY EFFICIENCY PRINCIPLE:
+    - Use agreement between solutions as evidence of correctness
+    - Use disagreement as a signal to investigate more deeply
+    - Only label a step as an error if it contains a specific mathematical mistake
+    - Incompleteness is not a mathematical error.
+
+    Here are the instructions for how to complete your task:
+
+    EFFICIENT VERIFICATION APPROACH:
+
+    1. QUICK COMPARISON (Use this to focus your attention):
+    - Immediately identify where the solutions differ in approach or results
+    - Use these differences as "error hotspots" to prioritize your verification
+    - When solutions agree, you can generally assume that part is correct
+    - When solutions disagree, investigate those specific points deeply
+
+    2. TARGETED VERIFICATION (Only where needed):
+    - Most important: Do not consider any incomplete steps as errors
+    - Focus your mathematical verification on the "hotspots" identified above
+    - Check mathematical validity only at points of difference or uncertainty
+    - Avoid line-by-line checking of steps where solutions agree
+    - For each potential error spot, verify if the mathematical reasoning is valid
+    - If an intermediate step is later corrected, do not penalize the solution for having the incorrect intermediate step
+
+    After your targeted verification, propose a score tuple (score_1, score_2):
+    - Score (1,1) if both partial solutions are valid
+    - Score (1,0) if only the first solution is valid
+    - Score (0,1) if only the second solution is valid
+    - Score (0,0) if both solutions are invalid
+
+    In case you score a solution as 0, you must give an explanation for each check below:
+    3. FINAL CHECKS:
+    - If you score a solution as 0, you MUST identify the specific mathematical error.
+    - You must also double check the problem statement. Reconsider your score and determine if you have misinterpreted the problem statement.
+    - You must also check whether you have penalized a solution for being incomplete or having missing steps.
+
+    Before outputting your final score, you must answer these questions:
+    STOP! Did you give a score of 0 to a solution that was incomplete?
+    STOP! Did you penalize a solution for being incomplete or having missing steps?
+    STOP! Did you make a comparison of which solution is better?
+    STOP! Did you consider steps incorrect just because they differ between solutions?
+    STOP! Did you prematurely evaluate based on final answers?
+    STOP! Did you consider a step incorrect just because it lacks sufficient detail or justification?
+    
+    Now give your final score:
+    Final score:
+    """
+    completion = client.chat.completions.create(
+        model=verification_model,
+        messages=[
+            {"role": "system", "content": "You are a helpful assistant."},
+            {"role": "user", "content": prompt},
+        ],
+        temperature=0.0,
+    )
+    response = completion.choices[0].message.content.strip()
+    print(full_text_1)
+    print(full_text_2)
+    print(f"Correctness score: {response} \n")
+    score_match = re.findall(r"\(\s*([01](?:\.0)?)\s*,\s*([01](?:\.0)?)\s*\)", response)
+    score = score_match[-1] if score_match else (0, 0)
+    return score
+
+
+def self_complete(verification_prompt: str, verification_model: str, api: str = "openai"):
+    print(verification_prompt)
+    """Completetion method"""
+    if api == "openai":
+        client = OpenAI()
+    elif api == "hf":
+        client = InferenceClient()
+    elif api == "together":
+        client = Together()
+    else:
+        raise ValueError(f"Invalid API: {api}")
+
+    completion = client.chat.completions.create(
+        model=verification_model,
+        messages=[
+            {"role": "system", "content": "You are a helpful assistant."},
+            {"role": "user", "content": verification_prompt},
+        ],
+        temperature=0.0,
+    )
+    response = completion.choices[0].message.content.strip()
+    return response

src/global_edit_utils.py

@@ -0,0 +1,127 @@
+from huggingface_hub import InferenceClient
+from openai import OpenAI
+
+bio_prompt = """
+You are given a piece of text that is a part of a biography of an entity. This text may contain some minor errors that make it incoherent as well as potentially redundant information. Your task is to fix the errors and make the text coherent.
+Then, remove any redundant information.
+Text: {text}
+
+If this is not possible because the text is just a fragment of a sentence, return "Abstain".
+If the text already claims a lack of knowledge about the topic, return "Abstain".
+Only return the cleaned up text. Do not include any other text:
+"""
+
+fp_prompt = """
+You are given a piece of text that is a part of a false presupposition task which includes outputting a list of items. 
+This text may contain some minor errors that make it incoherent as well as potentially redundant information. Your task is to fix the errors and make the text coherent.
+Then, remove any redundant information.
+Text: {text}
+
+The resulting list of items should be separated by semicolons with no other text.
+If this list it not possible to generate, return "Abstain".
+"""
+
+hist_prompt = """
+You are given a piece of text that is a part of a historical event task. This text may contain some minor errors that make it incoherent as well as potentially redundant information. Your task is to fix the errors and make the text coherent.
+Then, remove any redundant information.
+Text: {text}
+
+If this is not possible because the text is just a fragment of a sentence, return "Abstain".
+If the text already claims a lack of knowledge about the topic, return "Abstain".
+Only return the cleaned up text. Do not include any other text:
+"""
+
+refs_prompt = """
+You are given a piece of text that is a part of a reference task. This text may contain some minor errors that make it incoherent as well as potentially redundant information. Your task is to fix the errors and make the text coherent.
+Then, remove any redundant information.
+Text: {text}
+
+If this is not possible because the text is just a fragment of a sentence, return "Abstain".
+If the text already claims a lack of knowledge about the topic, return "Abstain".
+Only return the cleaned up text. Do not include any other text:
+"""
+
+gpqa_prompt = """
+You are given a piece of text that is a part of a graduate level question answering task. This text may contain some minor errors that make it incoherent as well as potentially redundant information. Your task is to fix the errors and make the text coherent.
+Then, remove any redundant information.
+Text: {text}
+Only return the cleaned up text. Do not include any other text:
+"""
+
+popqa_prompt =  """
+You are given a piece of text that is a part of a paragraph which details facts related to an entity. This text may contain some minor errors that make it incoherent as well as potentially redundant information. Your task is to fix the errors and make the text coherent.
+Then, remove any redundant information.
+Text: {text}
+
+If this is not possible because the text is just a fragment of a sentence, return "Abstain".
+If the text already claims a lack of knowledge about the topic, return "Abstain".
+Only return the cleaned up text. Do not include any other text:
+"""
+task_to_prompt = {
+    "bio": bio_prompt,
+    "fp": fp_prompt,
+    "hist": hist_prompt,
+    "refs": refs_prompt,
+    "gpqa": gpqa_prompt,
+    "popqa": popqa_prompt
+}
+
+'''
+Cleans up disfluencies in the draft response in consensus decoding. 
+
+Args:
+    text: The text to clean up.
+    api: The API to use for cleaning up the text.
+    task: The task : biography, false presupposition, historical event, reference, graduate question answering, paragraph question answering.
+    model: The model to use for cleaning up the text.
+
+Returns:
+    A string of the cleaned up text.
+'''
+
+def clean_up_text(text: str, api: str, task: str, model: str = "gpt-4.1-mini", **kwargs):
+    if api == "openai":
+        client = OpenAI()
+    elif api == "hf":
+        tokenizer = kwargs.get("tokenizer")
+        model = kwargs.get("hf_model")
+        
+        if tokenizer is None or model is None:
+            raise ValueError("For 'hf', both 'tokenizer' and 'model' must be provided.")
+        
+        clean_up_prompt = task_to_prompt[task].format(text=text)
+        
+        messages = [{"role": "user", "content": clean_up_prompt}]
+        input_ids = tokenizer.apply_chat_template(
+            messages,
+            add_generation_prompt=True,
+            return_tensors="pt"
+        ).to(model.device)
+        
+        terminators = [ tokenizer.eos_token_id, ]
+        outputs = model.generate(
+            input_ids,
+            max_new_tokens=500,
+            do_sample=False,  
+            pad_token_id=tokenizer.eos_token_id,
+            eos_token_id=terminators,
+        )
+        
+        return tokenizer.decode(
+            outputs[0][input_ids.shape[-1]:],
+            skip_special_tokens=True
+        ).strip()
+    else:
+        raise ValueError(f"Invalid API: {api}")
+
+    clean_up_prompt = task_to_prompt[task].format(text=text)
+
+    completion = client.chat.completions.create(
+        model=model,
+        messages=[
+            {"role": "system", "content": "You are a helpful assistant."},
+            {"role": "user", "content": clean_up_prompt},
+        ],
+    )
+
+    return completion.choices[0].message.content.strip()

src/new_alignment.py

@@ -0,0 +1,150 @@
+import numpy
+
+
+class ScoreParam:
+    def __init__(self, match, mismatch, gap_open, gap_extend):
+        self.match = match
+        self.mismatch = mismatch
+        self.gap_open = gap_open
+        self.gap_extend = gap_extend
+
+    def __str__(self):
+        return f"Match: {self.match}, Mismatch: {self.mismatch}, Gap Open: {self.gap_open}, Gap Extend: {self.gap_extend}"
+
+
+class SeqGraphAlignment(object):
+    __default_score = ScoreParam(1, -3, -2, -1)
+
+    def __init__(
+        self,
+        sequence,
+        graph,
+        fastMethod=True,
+        globalAlign=False,
+        score_params=__default_score,
+        *args,
+        **kwargs,
+    ):
+        self.score = score_params
+        self.sequence = sequence
+        self.graph = graph
+        self.stringidxs = None
+        self.nodeidxs = None
+        self.globalAlign = globalAlign
+        if fastMethod:
+            matches = self.alignStringToGraphFast(*args, **kwargs)
+        else:
+            matches = self.alignStringToGraphSimple(*args, **kwargs)
+        self.stringidxs, self.nodeidxs = matches
+
+    def alignmentStrings(self):
+        return (
+            "".join(self.sequence[i] if i is not None else "-" for i in self.stringidxs),
+            "".join(self.graph.nodedict[j].text if j is not None else "-" for j in self.nodeidxs),
+        )
+
+    def matchscore(self, c1, c2):
+        if c1 == c2:
+            return self.score.match
+        else:
+            return self.score.mismatch
+
+    def matchscoreVec(self, c, v):
+        return numpy.where(v == c, self.score.match, self.score.mismatch)
+
+    def prevIndices(self, node, nodeIDtoIndex):
+        prev = [nodeIDtoIndex[predID] for predID in list(node.inEdges.keys())]
+        if not prev:
+            prev = [-1]
+        return prev
+
+    def initializeDynamicProgrammingData(self):
+        l1 = self.graph.nNodes
+        l2 = len(self.sequence)
+
+        nodeIDtoIndex = {}
+        nodeIndexToID = {-1: None}
+        ni = self.graph.nodeiterator()
+        for index, node in enumerate(ni()):
+            nodeIDtoIndex[node.ID] = index
+            nodeIndexToID[index] = node.ID
+
+        scores = numpy.zeros((3, l1 + 1, l2 + 1), dtype=numpy.int32)
+
+        if self.globalAlign:
+            # M[0, i] = -inf
+            scores[0, 0, :] = [
+                -1000000000 for i in range(l2+1)
+            ]
+            scores[0, 0, 0] = 0
+            # X[0, i] = gap_open + i * gap_extend
+            scores[1, 0, :] = [
+                self.score.gap_open + i * self.score.gap_extend for i in range(l2 + 1) 
+            ]
+            scores[1, 0, 0] = -1000000000
+            # Y[0, i] = -inf
+            scores[2, 0, :] = [
+                -1000000000 for i in range(l2+1)
+            ]
+
+            ni = self.graph.nodeiterator()
+            # After topology sort, the predcessors will have index less than the current node
+            for index, node in enumerate(ni()):
+                scores[0, index + 1, 0] = -1000000000
+                scores[1, index + 1, 0] = -1000000000
+                prevIdxs = self.prevIndices(node, nodeIDtoIndex)
+                best = scores[2 ,prevIdxs[0] + 1, 0]
+                for prevIdx in prevIdxs:
+                    best = max(best, scores[2, prevIdx + 1, 0])
+                # If we have no predecessors, we start the gap 
+                if prevIdxs == [-1]:
+                    scores[2, index + 1, 0] =  self.score.gap_open + self.score.gap_extend
+                else:
+                    scores[2, index + 1, 0] = best + self.score.gap_extend
+
+        # 3D Backtracking
+        backStrIdx = numpy.zeros((3, l1 + 1, l2 + 1), dtype=numpy.int32)
+        backGrphIdx = numpy.zeros((3, l1 + 1, l2 + 1), dtype=numpy.int32)
+        backMtxIdx = numpy.zeros((3, l1 + 1, l2 + 1), dtype=numpy.int32)
+
+        return nodeIDtoIndex, nodeIndexToID, scores, backStrIdx, backGrphIdx, backMtxIdx
+
+    def backtrack(self, scores, backStrIdx, backGrphIdx, backMtxIdx ,nodeIndexToID):
+        besti, bestj = scores.shape[1] - 1, scores.shape[2] - 1
+        #Storing best matrices for each [i,j]
+        scores_arr = numpy.array(scores)
+        max_m = numpy.argmax(scores_arr, axis=0)
+        
+        if self.globalAlign:
+            ni = self.graph.nodeiterator()
+            # Finding the best node to start from
+            terminalIndices = [index for (index, node) in enumerate(ni()) if node.outDegree == 0]
+            print(terminalIndices)
+            besti = terminalIndices[0] + 1
+            bestscore = scores[max_m[besti, bestj], besti, bestj]
+            for i in terminalIndices[1:]:
+                score = scores[max_m[i + 1, bestj], i + 1, bestj]
+                if score > bestscore:
+                    bestscore, besti = score, i + 1
+            bestm = max_m[besti, bestj]
+
+        matches = []
+        strindexes = []
+
+        while (besti != 0 or bestj != 0):
+            nextm, nexti, nextj,  = backMtxIdx[bestm, besti, bestj], backGrphIdx[bestm, besti, bestj], backStrIdx[bestm, besti, bestj]
+            curstridx, curnodeidx = bestj - 1, nodeIndexToID[besti - 1]
+            
+            if bestm == 0:
+                matches.insert(0, curnodeidx)
+                strindexes.insert(0, curstridx)
+            elif bestm == 1:
+                matches.insert(0, None)
+                strindexes.insert(0, curstridx)
+            else:
+                matches.insert(0, curnodeidx)
+                strindexes.insert(0, None)
+
+            bestm, besti, bestj = nextm, nexti, nextj
+
+        return strindexes, matches

src/new_text_alignment.py

@@ -0,0 +1,134 @@
+from difflib import SequenceMatcher
+
+import numpy as np
+
+from .new_alignment import ScoreParam, SeqGraphAlignment
+
+PUNCTUATION_MARKS = [".", "!", "?", ",", ":", ";", "...", "(", ")"]
+
+class TextSeqGraphAlignment(SeqGraphAlignment):
+    def __init__(
+        self,
+        text,
+        graph,
+        fastMethod=True,
+        globalAlign=True,
+        matchscore=1,
+        mismatchscore=-3,
+        gap_open=-2,
+        gap_extend=-1,
+        position_weight=0.1,
+        *args,
+        **kwargs,
+    ):
+        score_params = ScoreParam(
+            match=matchscore, mismatch=mismatchscore, gap_open=gap_open, gap_extend=gap_extend
+        )
+
+        if isinstance(text, str):
+            self.original_text = text
+            self.sequence = text.split()
+        else:
+            self.sequence = text
+            self.original_text = " ".join(text)
+        self.position_weight = position_weight
+
+        super().__init__(
+            self.sequence,
+            graph,
+            fastMethod,
+            globalAlign=globalAlign,
+            score_params=score_params,
+            *args,
+            **kwargs,
+        )
+
+    def string_similarity(self, s1, s2):
+        """Get edit-distance based similarity between two strings"""
+        return SequenceMatcher(None, s1, s2).ratio()
+
+    def matchscore(self, word1: str, word2: str) -> float:
+        """Enhanced scoring function that considers string similarity
+        and relative position"""
+        # Calculate basic string similarity
+        similarity = self.string_similarity(word1, word2)
+
+        # If words are very similar, treat as match
+        if similarity > 0.8:  # Can tune this threshold
+            similarity = self.score.match
+        # For less similar words, scale score based on similarity
+        elif similarity > 0.5:  # Can tune this threshold too
+            similarity = self.score.match * similarity
+        else:
+            similarity = self.score.mismatch
+            return similarity
+
+        # add weight if any punctuation mark is present
+        if any(char in word1 for char in PUNCTUATION_MARKS) or any(
+            char in word2 for char in PUNCTUATION_MARKS
+        ):
+            similarity = similarity * 1.5
+
+        return similarity
+
+    def alignmentStrings(self):
+        """Override to handle word-based alignment"""
+        aligned_seq = [self.sequence[i] if i is not None else "-" for i in self.stringidxs]
+        aligned_graph = [
+            self.graph.nodedict[j].text if j is not None else "-" for j in self.nodeidxs
+        ]
+        return " ".join(aligned_seq), " ".join(aligned_graph)
+
+    def alignStringToGraphFast(self):
+        if not isinstance(self.sequence, list):
+            raise TypeError("Sequence must be a list of words")
+
+        nodeIDtoIndex, nodeIndexToID, scores, backStrIdx, backGrphIdx, backMtxIdx = (
+            self.initializeDynamicProgrammingData()
+        )
+        # M: Match at last indices, X: Gap at last index of graph, Y: gap at last index of sequence
+        M, X, Y = 0, 1, 2
+        
+        ni = self.graph.nodeiterator()
+        for i, node in enumerate(ni()):
+            gbase = node.text
+            
+            for j, sbase in enumerate(self.sequence):
+                candidates_X , candidates_Y , candidates_M = [], [], []
+                candidates_X += [
+                    (self.score.gap_open + self.score.gap_extend + scores[0, i + 1, j], i + 1, j, M),
+                    (self.score.gap_extend + scores[1, i + 1, j], i + 1, j, X),
+                    (self.score.gap_open + self.score.gap_extend + scores[2, i + 1, j], i + 1, j, Y)
+                ]
+                for predIndex in self.prevIndices(node, nodeIDtoIndex):
+                    candidates_Y += [
+                        (self.score.gap_open + self.score.gap_extend + scores[0, predIndex + 1, j + 1] , predIndex + 1, j + 1, M),
+                        (self.score.gap_open + self.score.gap_extend + scores[1, predIndex + 1, j + 1] , predIndex + 1, j + 1, X),
+                        (self.score.gap_extend + scores[2, predIndex + 1, j + 1] , predIndex + 1, j + 1, Y)
+                    ]
+                    candidates_M += [
+                        (self.matchscore(sbase, gbase) +  scores[0, predIndex + 1, j], predIndex + 1, j, M),
+                        (self.matchscore(sbase, gbase) +  scores[1, predIndex + 1, j], predIndex + 1, j, X),
+                        (self.matchscore(sbase, gbase) +  scores[2, predIndex + 1, j], predIndex + 1, j, Y)
+                    ]
+                
+                (
+                    scores[0, i + 1, j + 1],
+                    backGrphIdx[0, i + 1, j + 1],
+                    backStrIdx[0, i + 1, j + 1],
+                    backMtxIdx[0, i + 1, j + 1],
+                ) = max(candidates_M)
+                (
+                    scores[1, i + 1, j + 1],
+                    backGrphIdx[1, i + 1, j + 1],
+                    backStrIdx[1, i + 1, j + 1],
+                    backMtxIdx[1, i + 1, j + 1],
+                ) = max(candidates_X)
+                (
+                    scores[2, i + 1, j + 1],
+                    backGrphIdx[2, i + 1, j + 1],
+                    backStrIdx[2, i + 1, j + 1],
+                    backMtxIdx[2, i + 1, j + 1],
+                ) = max(candidates_Y)        
+                
+        return self.backtrack(scores, backStrIdx, backGrphIdx, backMtxIdx ,nodeIndexToID)

src/poa_graph.py

@@ -0,0 +1,685 @@
+"""
+Adapted from Jonathan Dursi
+https://github.com/ljdursi/poapy
+"""
+
+import collections
+import textwrap
+from typing import Dict, List, Optional, Union
+
+import numpy
+
+from .alignment import SeqGraphAlignment
+
+
+class Node(object):
+    def __init__(self, nodeID: int = -1, text: str = ""):
+        self.ID = nodeID
+        self.text = text
+        self.inEdges = {}
+        self.outEdges = {}
+        self.alignedTo = []
+
+    def __str__(self):
+        return "(%d:%s)" % (self.ID, self.text)
+
+    def _add_edge(
+        self,
+        edgeset: Dict[int, "Node"],
+        neighbourID: int,
+        label: Union[int, List[int]],
+        from_neighbour: bool,
+        weight: int = 1,
+    ):
+        if neighbourID is None:
+            return
+        # already present? just update labels
+        # otherwise create appropriately-ordered edge and proceed
+        if neighbourID in edgeset:
+            edgeset[neighbourID].weight += weight
+            if isinstance(label, list):
+                edgeset[neighbourID].labels.extend(label)
+            else:
+                edgeset[neighbourID].labels.append(label)
+            # remove duplicates
+            edgeset[neighbourID].labels = list(set(edgeset[neighbourID].labels))
+        else:
+            if from_neighbour:
+                edge = Edge(outNodeID=neighbourID, inNodeID=self.ID, label=label, weight=weight)
+            else:
+                edge = Edge(outNodeID=self.ID, inNodeID=neighbourID, label=label, weight=weight)
+            edgeset[neighbourID] = edge
+
+    def addInEdge(self, neighbourID: int, label: Optional[Union[int, List[int]]], weight: int = 1):
+        self._add_edge(self.inEdges, neighbourID, label, from_neighbour=True, weight=weight)
+
+    def addOutEdge(self, neighbourID: int, label: Optional[Union[int, List[int]]], weight: int = 1):
+        self._add_edge(self.outEdges, neighbourID, label, from_neighbour=False, weight=weight)
+
+    def nextNode(self, label: int):
+        """Returns the first (presumably only) outward neighbour
+        having the given edge label"""
+        nextID = None
+        for e in self.outEdges:
+            if label in self.outEdges[e].labels:
+                nextID = e
+        return nextID
+
+    @property
+    def inDegree(self):
+        return len(self.inEdges)
+
+    @property
+    def outDegree(self):
+        return len(self.outEdges)
+
+    @property
+    def weightedInDegree(self):
+        return sum(edge.weight for edge in self.inEdges.values())
+
+    @property
+    def weightedOutDegree(self):
+        return sum(edge.weight for edge in self.outEdges.values())
+
+    @property
+    def labels(self):
+        """Returns all the labels associated with an in-edge or an out edge."""
+        labelset = set([])
+        for e in list(self.inEdges.values()):
+            labelset = labelset.union(e.labels)
+        for e in list(self.outEdges.values()):
+            labelset = labelset.union(e.labels)
+        return list(labelset)
+
+
+class Edge(object):
+    def __init__(
+        self,
+        inNodeID: int = -1,
+        outNodeID: int = -1,
+        label: Optional[Union[int, List[int]]] = None,
+        weight: int = 1,
+    ):
+        self.inNodeID = inNodeID
+        self.outNodeID = outNodeID
+
+        self.weight = weight
+
+        if label is None:
+            self.labels = []
+        elif isinstance(label, list):
+            self.labels = label
+        else:
+            self.labels = [label]
+
+    def addLabel(self, newlabel):
+        self.labels.append(newlabel)
+
+    def __str__(self):
+        nodestr = "(%d) -> (%d) " % (self.inNodeID, self.outNodeID)
+        if self.labels is None:
+            return nodestr
+        else:
+            return nodestr + self.labels.__str__()
+
+
+class POAGraph(object):
+    def addUnmatchedSeq(self, seq, label: int = -1, updateSequences=True):
+        """Add a completely independant (sub)string to the graph,
+        and return node index to initial and final node"""
+        if seq is None:
+            return
+
+        firstID, lastID = None, None
+        neededSort = self.needsSort
+
+        for text in seq:
+            nodeID = self.addNode(text)
+            if firstID is None:
+                firstID = nodeID
+            if lastID is not None:
+                self.addEdge(lastID, nodeID, label)
+            lastID = nodeID
+
+        self._needsort = neededSort  # no new order problems introduced
+        if updateSequences:
+            self._seqs.append(seq)
+            self._labels.append(label)
+            self._starts.append(firstID)
+        return firstID, lastID
+
+    def __init__(self, seq=None, label: Optional[Union[int, List[int]]] = None):
+        self._nextnodeID = 0
+        self._nnodes = 0
+        self._nedges = 0
+        self.nodedict = {}
+        self.nodeidlist = []  # allows a (partial) order to be imposed on the nodes
+        self._needsort = False
+        self._labels = []
+        self._seqs = []
+        self._starts = []
+
+        if seq is not None:
+            self.addUnmatchedSeq(seq, label)
+
+    def nodeIdxToBase(self, idx):
+        return self.nodedict[self.nodeidlist[idx]].text
+
+    def addNode(self, text):
+        nid = self._nextnodeID
+        newnode = Node(nid, text)
+        self.nodedict[nid] = newnode
+        self.nodeidlist.append(nid)
+        self._nnodes += 1
+        self._nextnodeID += 1
+        self._needsSort = True
+        return nid
+
+    def addEdge(self, start, end, label, weight: int = 1):
+        if start is None or end is None:
+            return
+
+        if start not in self.nodedict:
+            raise KeyError("addEdge: Start node not in graph: " + str(start))
+        if end not in self.nodedict:
+            raise KeyError("addEdge: End node not in graph: " + str(end))
+
+        oldNodeEdges = self.nodedict[start].outDegree + self.nodedict[end].inDegree
+
+        self.nodedict[start].addOutEdge(end, label, weight)
+        self.nodedict[end].addInEdge(start, label, weight)
+
+        newNodeEdges = self.nodedict[start].outDegree + self.nodedict[end].inDegree
+
+        if newNodeEdges != oldNodeEdges:
+            self._nedges += 1
+
+        self._needsSort = True
+        return
+
+    @property
+    def needsSort(self):
+        return self._needsort
+
+    @property
+    def nNodes(self):
+        return self._nnodes
+
+    @property
+    def nEdges(self):
+        return self._nedges
+
+    @property
+    def num_sequences(self):
+        return len(self._seqs)
+
+    def get_sequences(self):
+        return self._seqs
+
+    def _simplified_graph_rep(self):
+
+        node_to_pn = {}
+        pn_to_nodes = {}
+
+        # Find the mappings from nodes to pseudonodes
+        cur_pnid = 0
+        for _, node in self.nodedict.items():
+            if node.ID not in node_to_pn:
+                node_ids = [node.ID] + node.alignedTo
+                pn_to_nodes[cur_pnid] = node_ids
+                for nid in node_ids:
+                    node_to_pn[nid] = cur_pnid
+                cur_pnid += 1
+
+        # create the pseudonodes
+        Pseudonode = collections.namedtuple(
+            "Pseudonode", ["pnode_id", "predecessors", "successors", "node_ids"]
+        )
+        pseudonodes = []
+
+        for pnid in range(cur_pnid):
+            nids, preds, succs = pn_to_nodes[pnid], [], []
+            for nid in nids:
+                node = self.nodedict[nid]
+                preds += [node_to_pn[inEdge.outNodeID] for _, inEdge in node.inEdges.items()]
+                succs += [node_to_pn[outEdge.inNodeID] for _, outEdge in node.outEdges.items()]
+
+            pn = Pseudonode(pnode_id=pnid, predecessors=preds, successors=succs, node_ids=nids)
+            pseudonodes.append(pn)
+
+        return pseudonodes
+
+    def toposort(self):
+        """Sorts node list so that all incoming edges come from nodes earlier in the list."""
+        sortedlist = []
+        completed = set([])
+
+        #
+        # The topological sort of this graph is complicated by the alignedTo edges;
+        # we want to nodes connected by such edges to remain near each other in the
+        # topological sort.
+        #
+        # Here we'll create a simple version of the graph that merges nodes that
+        # are alignedTo each other, performs the sort, and then decomposes the
+        # 'pseudonodes'.
+        #
+        # The need for this suggests that the way the graph is currently represented
+        # isn't quite right and needs some rethinking.
+        #
+
+        pseudonodes = self._simplified_graph_rep()
+
+        def dfs(start, complete, sortedlist):
+            stack, started = [start], set()
+            while stack:
+                pnodeID = stack.pop()
+
+                if pnodeID in complete:
+                    continue
+
+                if pnodeID in started:
+                    complete.add(pnodeID)
+                    for nid in pseudonodes[pnodeID].node_ids:
+                        sortedlist.insert(0, nid)
+                    started.remove(pnodeID)
+                    continue
+
+                successors = pseudonodes[pnodeID].successors
+                started.add(pnodeID)
+                stack.append(pnodeID)
+                stack.extend(successors)
+
+        while len(sortedlist) < self.nNodes:
+            found = None
+            for pnid in range(len(pseudonodes)):
+                if pnid not in completed and len(pseudonodes[pnid].predecessors) == 0:
+                    found = pnid
+                    break
+            assert found is not None
+            dfs(found, completed, sortedlist)
+
+        assert len(sortedlist) == self.nNodes
+        self.nodeidlist = sortedlist
+        self._needsSort = False
+        return
+
+    def testsort(self):
+        """Test the nodeidlist to make sure it is topologically sorted:
+        eg, all predecessors of a node preceed the node in the list"""
+        if self.nodeidlist is None:
+            return
+        seen_nodes = set()
+        for nodeidx in self.nodeidlist:
+            node = self.nodedict[nodeidx]
+            for in_neighbour in node.inEdges:
+                assert in_neighbour in seen_nodes
+            seen_nodes.add(nodeidx)
+        return
+
+    def nodeiterator(self):
+        if self.needsSort:
+            self.toposort()
+
+        def nodegenerator():
+            for nodeidx in self.nodeidlist:
+                yield self.nodedict[nodeidx]
+
+        return nodegenerator
+
+    def __str__(self):
+        selfstr = ""
+        ni = self.nodeiterator()
+        for node in ni():
+            selfstr += node.__str__() + "\n"
+            for outIdx in node.outEdges:
+                selfstr += "        " + node.outEdges[outIdx].__str__() + "\n"
+        return selfstr
+
+    def incorporateSeqAlignment(self, alignment: SeqGraphAlignment, seq, label: int = -1):
+        """Incorporate a SeqGraphAlignment into the graph."""
+        newseq = alignment.sequence
+        stringidxs = alignment.stringidxs
+        nodeidxs = alignment.nodeidxs
+
+        firstID = None
+        headID = None
+        tailID = None
+
+        path = []
+        # head, tail of sequence may be unaligned; just add those into the
+        # graph directly
+        validstringidxs = [si for si in stringidxs if si is not None]
+        startSeqIdx, endSeqIdx = validstringidxs[0], validstringidxs[-1]
+        if startSeqIdx > 0:
+            firstID, headID = self.addUnmatchedSeq(
+                newseq[0:startSeqIdx], label, updateSequences=False
+            )
+        if endSeqIdx < len(newseq):
+            tailID, __ = self.addUnmatchedSeq(newseq[endSeqIdx + 1 :], label, updateSequences=False)
+
+        # now we march along the aligned part. For each text, we find or create
+        # a node in the graph:
+        #   - if unmatched, the corresponding node is a new node
+        #   - if matched:
+        #       - if matched to a node with the same text, the node is that node
+        #       - if matched to a node with a different text whch is in turn
+        #         aligned to a node with the same text, that aligned node is
+        #         the node
+        #       - otherwise, we create a new node.
+        # In all cases, we create edges (or add labels) threading through the
+        # nodes.
+        for sindex, matchID in zip(stringidxs, nodeidxs):
+            if sindex is None:
+                continue
+            text = newseq[sindex]
+            if matchID is None:
+                nodeID = self.addNode(text)
+            elif self.nodedict[matchID].text == text:
+                nodeID = matchID
+            else:
+                otherAligns = self.nodedict[matchID].alignedTo
+                foundNode = None
+                for otherNodeID in otherAligns:
+                    if self.nodedict[otherNodeID].text == text:
+                        foundNode = otherNodeID
+                if foundNode is None:
+                    nodeID = self.addNode(text)
+                    self.nodedict[nodeID].alignedTo = [matchID] + otherAligns
+                    for otherNodeID in [matchID] + otherAligns:
+                        self.nodedict[otherNodeID].alignedTo.append(nodeID)
+                else:
+                    nodeID = foundNode
+
+            self.addEdge(headID, nodeID, label)
+            headID = nodeID
+            if firstID is None:
+                firstID = headID
+
+            path.append(nodeID)
+
+        # finished the unaligned portion: now add an edge from the current headID to the tailID.
+        self.addEdge(headID, tailID, label)
+
+        # resort
+        self.toposort()
+
+        self._seqs.append(seq)
+        self._labels.append(label)
+        self._starts.append(firstID)
+        self._seq_paths[label] = path
+        return
+
+    def consensus(self, excludeLabels=None):
+        if excludeLabels is None:
+            excludeLabels = []
+
+        if self.needsSort:
+            self.toposort()
+
+        nodesInReverse = self.nodeidlist[::-1]
+        maxnodeID = max(nodesInReverse) + 1
+        nextInPath = [-1] * maxnodeID
+        scores = numpy.zeros((maxnodeID))
+
+        for nodeID in nodesInReverse:
+            bestWeightScoreEdge = (-1, -1, None)
+            for neighbourID in self.nodedict[nodeID].outEdges:
+                # print(f"nodeID: {nodeID}, neighbourID: {neighbourID}")
+                e = self.nodedict[nodeID].outEdges[neighbourID]
+                weightScoreEdge = (e.weight, scores[neighbourID], neighbourID)
+
+                if weightScoreEdge > bestWeightScoreEdge:
+                    bestWeightScoreEdge = weightScoreEdge
+
+            scores[nodeID] = sum(bestWeightScoreEdge[0:2])
+            nextInPath[nodeID] = bestWeightScoreEdge[2]
+
+        pos = numpy.argmax(scores)
+        path = []
+        bases = []
+        labels = []
+        while pos is not None and pos > -1:
+            path.append(pos)
+            bases.append(self.nodedict[pos].text)
+            labels.append(self.nodedict[pos].labels)
+            pos = nextInPath[pos]
+
+        # ignore END node
+        path = path[:-1]
+        bases = bases[:-1]
+        labels = labels[:-1]
+        return path, bases, labels
+
+    def allConsenses(self, maxfraction=0.5):
+        allpaths = []
+        allbases = []
+        alllabels = []
+        exclusions = []
+
+        passno = 0
+        lastlen = 1000
+        maxpasses = 10
+
+        while len(exclusions) < len(self._labels) and lastlen >= 10 and passno < maxpasses:
+            path, bases, labellists = self.consensus(exclusions)
+            if len(path) > 0:
+                allpaths.append(path)
+                allbases.append(bases)
+                alllabels.append(labellists)
+
+                labelcounts = collections.defaultdict(int)
+                for ll in labellists:
+                    for label in ll:
+                        labelcounts[label] += 1
+
+                for label, seq in zip(self._labels, self._seqs):
+                    if label in labelcounts and labelcounts[label] >= maxfraction * len(seq):
+                        exclusions.append(label)
+
+            lastlen = len(path)
+            passno += 1
+
+        return list(zip(allpaths, allbases, alllabels))
+
+    def generateAlignmentStrings(self):
+        """Return a list of strings corresponding to the alignments in the graph"""
+
+        # Step 1: assign node IDs to columns in the output
+        #  column_index[node.ID] is the position in the toposorted node list
+        #    of the node itself, or the earliest node it is aligned to.
+        column_index = {}
+        current_column = 0
+
+        # go through nodes in toposort order
+        ni = self.nodeiterator()
+        for node in ni():
+            other_columns = [
+                column_index[other] for other in node.alignedTo if other in column_index
+            ]
+            if other_columns:
+                found_idx = min(other_columns)
+            else:
+                found_idx = current_column
+                current_column += 1
+
+            column_index[node.ID] = found_idx
+
+        ncolumns = current_column
+
+        # Step 2: given the column indexes, populate the strings
+        #   corresponding to the sequences inserted in the graph
+        seqnames = []
+        alignstrings = []
+        for label, start in zip(self._labels, self._starts):
+            seqnames.append(label)
+            curnode_id = start
+            charlist = ["-"] * ncolumns
+            while curnode_id is not None:
+                node = self.nodedict[curnode_id]
+                charlist[column_index[curnode_id]] = node.text
+                curnode_id = node.nextNode(label)
+            alignstrings.append("".join(charlist))
+
+        # Step 3: Same as step 2, but with consensus sequences
+        consenses = self.allConsenses()
+        for i, consensus in enumerate(consenses):
+            seqnames.append("Consensus" + str(i))
+            charlist = ["-"] * ncolumns
+            for path, text in zip(consensus[0], consensus[1]):
+                charlist[column_index[path]] = text
+            alignstrings.append("".join(charlist))
+
+        return list(zip(seqnames, alignstrings))
+
+    def jsOutput(self, verbose: bool = False, annotate_consensus: bool = True):
+        """returns a list of strings containing a a description of the graph for viz.js, http://visjs.org"""
+
+        # get the consensus sequence, which we'll use as the "spine" of the
+        # graph
+        pathdict = {}
+        if annotate_consensus:
+            path, __, __ = self.consensus()
+        lines = ["var nodes = ["]
+
+        ni = self.nodeiterator()
+        count = 0
+        for node in ni():
+            line = "    {id:" + str(node.ID) + ', label: "' + str(node.ID) + ": " + node.text + '"'
+            if node.ID in pathdict and count % 5 == 0 and annotate_consensus:
+                line += (
+                    ", x: "
+                    + str(pathdict[node.ID])
+                    + ", y: 0 , fixed: { x:true, y:false},"
+                    + "color: '#7BE141', is_consensus:true},"
+                )
+            else:
+                line += "},"
+            lines.append(line)
+
+        lines[-1] = lines[-1][:-1]
+        lines.append("];")
+
+        lines.append(" ")
+
+        lines.append("var edges = [")
+        ni = self.nodeiterator()
+        for node in ni():
+            nodeID = str(node.ID)
+            for edge in node.outEdges:
+                target = str(edge)
+                weight = str(len(node.outEdges[edge].labels) + 1.5)
+                lines.append(
+                    "    {from: "
+                    + nodeID
+                    + ", to: "
+                    + target
+                    + ", value: "
+                    + weight
+                    + ", color: '#4b72b0', arrows: 'to'},"
+                )
+            if verbose:
+                for alignededge in node.alignedTo:
+                    # These edges indicate alignment to different bases, and are
+                    # undirected; thus make sure we only plot them once:
+                    if node.ID > alignededge:
+                        continue
+                    target = str(alignededge)
+                    lines.append(
+                        "    {from: "
+                        + nodeID
+                        + ", to: "
+                        + target
+                        + ', value: 1, style: "dash-line", color: "red"},'
+                    )
+
+        lines[-1] = lines[-1][:-1]
+        lines.append("];")
+        return lines
+
+    def htmlOutput(self, outfile, verbose: bool = False, annotate_consensus: bool = True):
+        header = """
+                  <!doctype html>
+                  <html>
+                  <head>
+                    <title>POA Graph Alignment</title>
+
+                    <script type="text/javascript" src="https://unpkg.com/vis-network@9.0.4/standalone/umd/vis-network.min.js"></script>
+                  </head>
+
+                  <body>
+
+                  <div id="loadingProgress">0%</div>
+
+                  <div id="mynetwork"></div>
+
+                  <script type="text/javascript">
+                    // create a network
+                  """
+        outfile.write(textwrap.dedent(header[1:]))
+        lines = self.jsOutput(verbose=verbose, annotate_consensus=annotate_consensus)
+        for line in lines:
+            outfile.write(line + "\n")
+        footer = """
+                  var container = document.getElementById('mynetwork');
+                  var data= {
+                    nodes: nodes,
+                    edges: edges,
+                  };
+                  var options = {
+                    width: '100%',
+                    height: '800px',
+                    physics: {
+                        enabled: false,
+                        stabilization: {
+                            updateInterval: 10,
+                        },
+                        hierarchicalRepulsion: {
+                            avoidOverlap: 0.9,
+                        },
+                    },
+                    edges: {
+                        color: {
+                            inherit: false
+                        }
+                    },
+                    layout: {
+                        hierarchical: {
+                            direction: "UD",
+                            sortMethod: "directed",
+                            shakeTowards: "roots",
+                            levelSeparation: 150, // Adjust as needed
+                            nodeSpacing: 100, // Adjust as needed
+                            treeSpacing: 200, // Adjust as needed
+                            parentCentralization: true,
+                        }
+                    }
+                  };
+                  var network = new vis.Network(container, data, options);
+                  
+                  network.on('beforeDrawing', function(ctx) {
+                    nodes.forEach(function(node) {
+                        if (node.isConsensus) {
+                            // Set the level of spine nodes to the bottom
+                            network.body.data.nodes.update({
+                                id: node.id,
+                                level: 0 // Set level to 0 for spine nodes
+                            });
+                        }
+                    });
+                });
+
+                  network.on("stabilizationProgress", function (params) {
+                    document.getElementById("loadingProgress").innerText = Math.round(params.iterations / params.total * 100) + "%";
+                  });
+                  network.once("stabilizationIterationsDone", function () {
+                      document.getElementById("loadingProgress").innerText = "100%";
+                      setTimeout(function () {
+                        document.getElementById("loadingProgress").style.display = "none";
+                      }, 500);
+                  });
+                </script>
+
+                </body>
+                </html>
+                """
+        outfile.write(textwrap.dedent(footer))

src/text_poa_graph.py

@@ -0,0 +1,802 @@
+"""
+Enhanced version of POAGraph for text alignment
+"""
+
+import pickle
+import textwrap
+from typing import Dict, Optional
+
+import numpy as np
+from tqdm import tqdm
+
+from src.text_poa_graph_utils import path_sim_llm
+from src.global_edit_utils import clean_up_text
+
+from .new_text_alignment import TextSeqGraphAlignment
+from .poa_graph import Node, POAGraph
+
+
+class TextNode(Node):
+    def __init__(self, nodeID=-1, text=""):
+        super().__init__(nodeID, text)
+        self.variations = {}  # Track alternate phrasings
+        self.sequences = []  # Track sequences that contain this node
+        self.influenceScore = 0
+        self.num_tokens_used = 0
+
+    def add_variation(self, text, sequence_id):
+        self.variations[sequence_id] = text
+
+    @property
+    def is_stable(self):
+        """A node is stable if it appears frequently enough relative to total sequences"""
+        return self.frequency >= self.graph.stability_threshold
+
+
+class TextPOAGraph(POAGraph):
+    def __init__(self, text=None, label=-1):
+        self.consensus_node_ids = []
+        self._seq_paths = {}
+        self.end_id = -1
+        self.start_id = -1
+        self.failed = False
+        self.num_input_tokens_used = 0
+        self.num_output_tokens_used = 0
+        super().__init__(text, label)
+
+    def addNode(self, text):
+        """Override to use TextNode"""
+        nid = self._nextnodeID
+        newnode = TextNode(nid, text)
+        self.nodedict[nid] = newnode
+        self.nodeidlist.append(nid)
+        self._nnodes += 1
+        self._nextnodeID += 1
+        self._needsSort = True
+        return nid
+
+    def addUnmatchedSeq(self, text, label=-1, updateSequences=True):
+        """Modified to handle text sequences"""
+        if text is None:
+            return
+
+        # Handle both string and list input
+        if isinstance(text, str):
+            words = text.split()
+        else:
+            words = text
+
+        firstID, lastID = None, None
+        neededSort = self.needsSort
+
+        path = []
+        for word in words:
+            nodeID = self.addNode(word)
+            if firstID is None:
+                firstID = nodeID
+            if lastID is not None:
+                self.addEdge(lastID, nodeID, label=label)
+            lastID = nodeID
+            path.append(nodeID)
+
+        self._needsort = neededSort
+        if updateSequences:
+            self._seqs.append(words)
+            self._labels.append(label)
+            self._starts.append(firstID)
+            self._seq_paths[label] = path
+
+        return firstID, lastID
+
+    def add_text(self, text, label=-1):
+        """Main method to add new text to the alignment"""
+        if len(self._seqs) == 0:
+            # First sequence - just add it
+            self.addUnmatchedSeq(text, label)
+        else:
+            # Align to existing graph
+            alignment = TextSeqGraphAlignment(
+                text, self, matchscore=2, mismatchscore=-1, gapscore=-2
+            )
+            self.incorporateSeqAlignment(alignment, text, label)
+
+        # Update node frequencies
+        self._update_frequencies()
+
+    def removeNode(self, nodeID):
+        """Override to handle text nodes"""
+        node = self.nodedict[nodeID]
+        if node is None:
+            return
+
+        # Remove all edges to this node
+        out_edges = node.outEdges.copy()
+        in_edges = node.inEdges.copy()
+
+        for edge in out_edges:
+            self.removeEdge(node.ID, edge)
+        for edge in in_edges:
+            self.removeEdge(edge, node.ID)
+
+        # Remove from graph
+        del self.nodedict[nodeID]
+        self.nodeidlist.remove(nodeID)
+
+        for path in self._seq_paths.values():
+            if nodeID in path:
+                path.remove(nodeID)
+
+        self._nnodes -= 1
+        self._needsSort = True
+
+    def removeEdge(self, nodeID1, nodeID2):
+        """Override to handle text nodes"""
+        node1 = self.nodedict[nodeID1]
+        node2 = self.nodedict[nodeID2]
+
+        if node1 is None or node2 is None:
+            return
+
+        # Remove from graph
+        del node1.outEdges[nodeID2]
+        del node2.inEdges[nodeID1]
+
+    def merge_consensus_nodes(self, verbose: bool = False):
+        self.toposort()
+        # reset consensus node ids
+        self.consensus_node_ids = []
+        nodes = list(self.nodeiterator()())
+        consensus_segments = []
+        i = 0
+        while i < len(nodes):
+            node = nodes[i]
+            out_weight = sum(e.weight for e in node.outEdges.values())
+            in_weight = sum(e.weight for e in node.inEdges.values())
+
+            if out_weight in [0, self.num_sequences] and in_weight in [0, self.num_sequences]:
+                consensus_segment = [(node.ID, node.text)]
+                next_node = node
+                while (i + 1) < len(nodes) and len(next_node.outEdges) == 1:
+                    next_node = nodes[i + 1]
+                    next_out_weight = sum(e.weight for e in next_node.outEdges.values())
+                    next_in_weight = sum(e.weight for e in next_node.inEdges.values())
+
+                    if (
+                        next_out_weight != self.num_sequences
+                        or next_in_weight != self.num_sequences
+                    ):
+                        break
+
+                    consensus_segment.append((next_node.ID, next_node.text))
+                    i += 1
+                consensus_segments.append(consensus_segment)
+            i += 1
+        # merge consensus nodes into a single node
+        for segment in consensus_segments:
+            if len(segment) == 1:
+                self.consensus_node_ids.append(segment[0][0])
+                continue
+            merged_text = " ".join([text for _, text in segment])
+            first_node_id = segment[0][0]
+            last_node_id = segment[-1][0]
+
+            self.nodedict[last_node_id].text = merged_text
+            self.consensus_node_ids.append(last_node_id)
+
+            # attach all incoming edges to first node to last node
+            for id, edge in self.nodedict[first_node_id].inEdges.items():
+                weight = edge.weight
+                for _ in range(weight):
+                    self.addEdge(id, last_node_id, label=edge.labels)
+
+            # delete all nodes except last node
+            for node_id, _ in segment[:-1]:
+                self.removeNode(node_id)
+
+            
+
+        if verbose:
+            print(self.consensus_node_ids)
+
+    """
+    find all paths between start_node_id and end_node_id from original sequences
+    return a list of dictionaries with the following keys:
+    - path: list of node ids in the path (excluding start and including end)
+    - text: text of the path (excluding start and end)
+    - weight: minimal edge weight across all edges in the path
+    - labels: intersection of all edge labels in the path
+    """
+
+    def find_paths_between(self, start_node_id: int, end_node_id: int):
+        # find all paths between start_node_id and end_node_id from original sequences
+        path_dicts = []
+
+        # keep track of visited paths to avoid duplicates
+        visited_paths = set()
+
+        for _, path in self._seq_paths.items():
+            start_index = path.index(start_node_id) if start_node_id in path else None
+            end_index = path.index(end_node_id) if end_node_id in path else None
+
+            # print(start_index, end_index)
+            # print(path)
+
+            if (
+                start_index is not None
+                and end_index is not None
+                and end_index - start_index > 1
+                and tuple(path[start_index + 1 : end_index + 1]) not in visited_paths
+            ):
+                # intersection of all edge labels in the path
+                path_labels = set.intersection(
+                    *[
+                        set(self.nodedict[next_node_id].inEdges[node_id].labels)
+                        for node_id, next_node_id in zip(
+                            path[start_index:end_index], path[start_index + 1 : end_index + 1]
+                        )
+                    ]
+                )
+                path_weight = len(path_labels)
+                path_dicts.append(
+                    {
+                        
+                        "path": path[start_index + 1 : end_index + 1],
+                        "body_text": " ".join(
+                            [
+                                self.nodedict[node_id].text
+                                for node_id in path[start_index + 1 : end_index]
+                            ]
+                        ),
+                        "begin_text": self.nodedict[path[start_index]].text,
+                        "end_text": self.nodedict[path[end_index]].text,
+                        "weight": path_weight,
+                        "labels": path_labels,
+                    }
+                )
+                visited_paths.add(tuple(path[start_index + 1 : end_index + 1]))
+
+        return path_dicts
+
+    def _follow_path(self, start_id):
+        """Follow all possible paths from a node"""
+        paths = []
+        visited = set()
+
+        def dfs(node_id, current_path):
+            if node_id in visited:
+                return
+            visited.add(node_id)
+            node = self.nodedict[node_id]
+
+            if not node.outEdges:
+                paths.append(current_path + [node_id])
+                return
+
+            for next_id in node.outEdges:
+                dfs(next_id, current_path + [node_id])
+
+        dfs(start_id, [])
+        return paths
+
+    def merge_paths_between(
+        self,
+        start_node_id: int,
+        end_node_id: int,
+        path_sim_type: str = "llm",
+        verbose: bool = False,
+        **kwargs,
+    ):
+        path_dicts = self.find_paths_between(start_node_id, end_node_id)
+
+        if path_sim_type == "llm":
+            api = kwargs.get("api", "openai")
+            model = kwargs.get("model", "gpt-4o-mini")
+            domain = kwargs.get("domain", None)
+            similarity_judge_prompt = kwargs.get("similarity_judge_prompt", None)
+
+            def path_sim_func(path1_text, path2_text):
+                return path_sim_llm(
+                    path1_text,
+                    path2_text,
+                    api=api,
+                    model=model,
+                    domain=domain,
+                    custom_similarity_judge_prompt=similarity_judge_prompt,
+                )
+
+        elif path_sim_type == "cosine":
+            pass
+            # embedding_model = SentenceTransformer("all-MiniLM-L6-v2")
+            # threshold = kwargs.get("threshold", 0.9)
+            # path_sim_func = path_sim_cosine(embedding_model, threshold)
+        else:
+            raise ValueError(f"Invalid path similarity type: {path_sim_type}")
+
+        # merge paths based on semantic similarity
+        path_equivalence_classes = {}
+        class_count = 0
+
+        for path_dict in path_dicts:
+            if verbose:
+                print(path_dict)
+            found_class = False
+            for _, eq_class in path_equivalence_classes.items():
+                # check if path dict is already in an equivalence class
+                path1_text = (
+                    path_dict["begin_text"]
+                    + " "
+                    + path_dict["body_text"]
+                    + " "
+                    + path_dict["end_text"]
+                )
+                path2_text = (
+                    eq_class[0]["begin_text"]
+                    + " "
+                    + eq_class[0]["body_text"]
+                    + " "
+                    + eq_class[0]["end_text"]
+                )
+
+                judgement, num_input_tokens, num_output_tokens = path_sim_func(
+                    path1_text, path2_text
+                )
+                self.num_input_tokens_used += num_input_tokens
+                self.num_output_tokens_used += num_output_tokens
+                if judgement:
+                    eq_class.append(path_dict)
+                    found_class = True
+                    break
+            if not found_class:
+                class_count += 1
+                path_equivalence_classes[class_count] = [path_dict]
+
+        nodes_to_remove = set()  # Track nodes to remove
+        for _, eq_class in path_equivalence_classes.items():
+            path_dict = eq_class[0]
+
+            if verbose:
+                print(eq_class)
+            # add new node with merged text
+            new_node_id = self.addNode(path_dict["body_text"])
+            for sequence_id in path_dict["labels"]:
+                self.nodedict[new_node_id].variations[sequence_id] = path_dict["body_text"]
+
+            # collect nodes to remove from first path
+            nodes_to_remove.update(path_dict["path"][:-1])
+
+            # process data regarding weights and labels
+            labels = list(path_dict["labels"])
+            weight = path_dict["weight"]
+            self.addEdge(start_node_id, new_node_id, label=labels, weight=weight)
+
+            # Updated seq_paths for all labels to include new_node betwwen start_node and end_node
+            for label in labels:
+                index = self._seq_paths[label].index(start_node_id)
+                if (
+                    index + 1 < len(self._seq_paths[label])
+                    and self._seq_paths[label][index + 1] != new_node_id
+                ):
+                    self._seq_paths[label].insert(index + 1, new_node_id)
+
+            self.addEdge(new_node_id, end_node_id, label=labels, weight=weight)
+
+            self.nodedict[new_node_id].sequences = labels
+            # process additional paths
+            for path_dict in eq_class[1:]:
+                for sequence_id in path_dict["labels"]:
+                    self.nodedict[new_node_id].variations[sequence_id] = path_dict["body_text"]
+                nodes_to_remove.update(path_dict["path"][:-1])
+
+                # copy incoming edges to new node
+                labels = list(path_dict["labels"])
+                weight = path_dict["weight"]
+                self.addEdge(start_node_id, new_node_id, label=labels, weight=weight)
+
+                # Updated seq_paths for all labels to include new_node betwwen start_node and end_node
+                for label in labels:
+                    index = self._seq_paths[label].index(start_node_id)
+                    if (
+                        index + 1 < len(self._seq_paths[label])
+                        and self._seq_paths[label][index + 1] != new_node_id
+                    ):
+                        self._seq_paths[label].insert(index + 1, new_node_id)
+
+                self.addEdge(new_node_id, end_node_id, label=labels, weight=weight)
+                self.nodedict[new_node_id].sequences.extend(labels)
+
+            self.nodedict[new_node_id].sequences = list(set(self.nodedict[new_node_id].sequences))
+
+        # Remove all collected nodes after processing
+        for node_id in nodes_to_remove:
+            if node_id in self.nodedict:
+                if verbose:
+                    print(f"Removing node {node_id}")
+                self.removeNode(node_id)
+
+    def merge_divergent_paths(self, path_sim_type: str = "llm", verbose: bool = False, **kwargs):
+        # add dummy end node to the end of the graph
+        if not self.consensus_node_ids:
+            self.merge_consensus_nodes(verbose=verbose)
+
+        self.toposort()
+
+        if self.start_id == -1:
+            if verbose:
+                print("Adding start node")
+            self.start_id = self.addNode(text="START")
+            self._nextnodeID += 1
+            self.consensus_node_ids.insert(0, self.start_id)
+
+            for label, path in self._seq_paths.items():
+                self.addEdge(self.start_id, path[0], label=label, weight=1)
+                path.insert(0, self.start_id)
+
+        if self.end_id == -1:
+            if verbose:
+                print("Adding end node")
+            self.end_id = self.addNode(text="END")
+            self._nextnodeID += 1
+            self.consensus_node_ids = self.consensus_node_ids + [self.end_id]
+
+            for label, path in self._seq_paths.items():
+                self.addEdge(path[-1], self.end_id, label=label, weight=1)
+                path.append(self.end_id)
+
+        for i in tqdm(range(len(self.consensus_node_ids) - 1)):
+            if verbose:
+                print(self.consensus_node_ids[i], self.consensus_node_ids[i + 1])
+            self.merge_paths_between(
+                self.consensus_node_ids[i],
+                self.consensus_node_ids[i + 1],
+                path_sim_type=path_sim_type,
+                verbose=verbose,
+                **kwargs,
+            )
+
+    def get_variable_node_ids(self):
+        return [
+            node.ID for node in self.nodedict.values() if node.ID not in self.consensus_node_ids
+        ]
+
+    def compress_paths_between(self, start_node_id: int, end_node_id: int):
+        pass
+
+    def compress_graph(self):
+        pass
+
+    def update_influence_scores(self, outcome: Dict[int, float], discount_factor: float = 0.2):
+        self.toposort()
+        direct_scores = []
+        for node in self.nodedict.values():
+            next_out_weight = sum(e.weight for e in node.outEdges.values())
+            next_in_weight = sum(e.weight for e in node.inEdges.values())
+            if next_out_weight == self.num_sequences and next_in_weight == self.num_sequences:
+                out_list = []
+                for edge in node.outEdges.values():
+                    for _ in range(len(set(edge.labels))):
+                        out_list.append(np.mean([outcome[label] for label in set(edge.labels)]))
+                direct_scores.append((node.ID, np.var(out_list)))
+
+        scores = direct_scores.copy()
+
+        # Start from the end and propagate influence backward
+        for i in range(len(scores) - 2, -1, -1):
+            # Current node gets its direct influence plus discounted influence of next node
+            current_direct = scores[i][1]
+            next_total = scores[i + 1][1]
+            scores[i] = (scores[i][0], current_direct + discount_factor * next_total)
+
+        scores.sort(key=lambda x: x[1], reverse=True)
+        return scores
+
+    def jsOutput(
+        self,
+        verbose: bool = False,
+        annotate_consensus: bool = True,
+        color_annotations: Dict[int, str] = None,
+    ):
+        """returns a list of strings containing a a description of the graph for viz.js, http://visjs.org"""
+
+        # get the consensus sequence, which we'll use as the "spine" of the
+        # graph
+        pathdict = {}
+        if annotate_consensus:
+            path, __, __ = self.consensus()
+        lines = ["var nodes = ["]
+
+        ni = self.nodeiterator()
+        count = 0
+        for node in ni():
+            title_text = ""
+            if node.sequences:
+                title_text += f"Sequences: {node.sequences}"
+            if node.variations:
+                title_text += ";;;".join(
+                    [f"{sequence_id}: {text}" for sequence_id, text in node.variations.items()]
+                )
+                title_text = title_text.replace('"', "'")
+            line = (
+                "    {id:"
+                + str(node.ID)
+                + ', label: "'
+                + str(node.ID)
+                + ": "
+                + node.text.replace('"', "'")
+                + '", title: '
+                + '"'
+                + title_text
+                + '",'
+            )
+            if color_annotations and node.ID in color_annotations:
+                line += f" color: '{color_annotations[node.ID]}', "
+            if node.ID in pathdict and count % 5 == 0 and annotate_consensus:
+                line += (
+                    ", x: "
+                    + str(pathdict[node.ID])
+                    + ", y: 0 , fixed: { x:true, y:false},"
+                    + "color: '#7BE141', is_consensus:true},"
+                )
+            else:
+                line += "},"
+            lines.append(line)
+
+        lines[-1] = lines[-1][:-1]
+        lines.append("];")
+
+        lines.append(" ")
+
+        lines.append("var edges = [ ")
+        ni = self.nodeiterator()
+        for node in ni():
+            nodeID = str(node.ID)
+            for edge in node.outEdges:
+                target = str(edge)
+                weight = str(node.outEdges[edge].weight + 1.5)
+                lines.append(
+                    "    {from: "
+                    + nodeID
+                    + ", to: "
+                    + target
+                    + ", value: "
+                    + weight
+                    + ", color: '#4b72b0', arrows: 'to'},"
+                )
+            if verbose:
+                for alignededge in node.alignedTo:
+                    # These edges indicate alignment to different bases, and are
+                    # undirected; thus make sure we only plot them once:
+                    if node.ID > alignededge:
+                        continue
+                    target = str(alignededge)
+                    lines.append(
+                        "    {from: "
+                        + nodeID
+                        + ", to: "
+                        + target
+                        + ', value: 1, style: "dash-line", color: "red"},'
+                    )
+
+        lines[-1] = lines[-1][:-1]
+        lines.append("];")
+        return lines
+
+    def htmlOutput(
+        self,
+        outfile,
+        verbose: bool = False,
+        annotate_consensus: bool = True,
+        color_annotations: Dict[int, str] = None,
+    ):
+        header = """
+                  <!doctype html>
+                  <html>
+                  <head>
+                    <title>POA Graph Alignment</title>
+
+                    <script type="text/javascript" src="https://unpkg.com/vis-network@9.0.4/standalone/umd/vis-network.min.js"></script>
+                  </head>
+
+                  <body>
+
+                  <div id="loadingProgress">0%</div>
+
+                  <div id="mynetwork"></div>
+
+                  <script type="text/javascript">
+                    // create a network
+                  """
+        outfile.write(textwrap.dedent(header[1:]))
+        lines = self.jsOutput(
+            verbose=verbose,
+            annotate_consensus=annotate_consensus,
+            color_annotations=color_annotations,
+        )
+        for line in lines:
+            outfile.write(line + "\n")
+        footer = """
+                  var container = document.getElementById('mynetwork');
+                  var data= {
+                    nodes: nodes,
+                    edges: edges,
+                  };
+                  var options = {
+                    width: '100%',
+                    height: '800px',
+                    physics: {
+                        enabled: false,
+                        stabilization: {
+                            updateInterval: 10,
+                        },
+                    },
+                    edges: {
+                        color: {
+                            inherit: false
+                        }
+                    },
+                    layout: {
+                        hierarchical: {
+                            direction: "UD",
+                            sortMethod: "directed",
+                            shakeTowards: "roots",
+                            levelSeparation: 150, // Adjust as needed
+                            nodeSpacing: 800, // Adjust as needed
+                            treeSpacing: 200, // Adjust as needed
+                            parentCentralization: true,
+                        }
+                    }
+                  };
+                  var network = new vis.Network(container, data, options);
+                  
+                  network.on('beforeDrawing', function(ctx) {
+                    nodes.forEach(function(node) {
+                        if (node.isConsensus) {
+                            // Set the level of spine nodes to the bottom
+                            network.body.data.nodes.update({
+                                id: node.id,
+                                level: 0 // Set level to 0 for spine nodes
+                            });
+                        }
+                    });
+                });
+
+                  network.on("stabilizationProgress", function (params) {
+                    document.getElementById("loadingProgress").innerText = Math.round(params.iterations / params.total * 100) + "%";
+                  });
+                  network.once("stabilizationIterationsDone", function () {
+                      document.getElementById("loadingProgress").innerText = "100%";
+                      setTimeout(function () {
+                        document.getElementById("loadingProgress").style.display = "none";
+                      }, 500);
+                  });
+
+                
+                </script>
+
+                </body>
+                </html>
+                """
+        outfile.write(textwrap.dedent(footer))
+
+    
+    def multi_consensus_response(self, abstention_threshold: Optional[float] = None, filter: bool = True):
+        self.toposort()
+        nodesInReverse = self.nodeidlist[::-1]
+        maxnodeID = self.end_id
+        nextInPath = [-1] * maxnodeID
+        scores = np.zeros(len(self.nodeidlist))
+
+        id_to_index = {node_id: index for index, node_id in enumerate(self.nodeidlist)}
+        index_to_id = {index: node_id for index, node_id in enumerate(self.nodeidlist)}
+
+        for nodeID in nodesInReverse:
+            bestWeightScoreEdges = [(-1, -1, None)]
+            for neighbourID in self.nodedict[nodeID].outEdges:
+                # print(f"nodeID: {nodeID}, neighbourID: {neighbourID}")
+                e = self.nodedict[nodeID].outEdges[neighbourID]
+                weightScoreEdge = (e.weight, scores[id_to_index[neighbourID]], neighbourID)
+                
+                
+                if weightScoreEdge > bestWeightScoreEdges[0]:
+                    bestWeightScoreEdges = [weightScoreEdge]
+                elif weightScoreEdge == bestWeightScoreEdges[0] and filter:
+                    bestWeightScoreEdges.append(weightScoreEdge)
+
+            
+            scores[id_to_index[nodeID]] = sum(bestWeightScoreEdges[0][0:2])
+            if bestWeightScoreEdges[0][2] is not None:
+                nextInPath[id_to_index[nodeID]] = id_to_index[bestWeightScoreEdges[0][2]]
+            else:
+                nextInPath[id_to_index[nodeID]] = None
+
+        pos = np.argmax(scores)
+        path = []
+        text = []
+        labels = []
+
+        while pos is not None and pos > -1:
+            if abstention_threshold is not None and self.nodedict[index_to_id[pos]].variations:
+                if (
+                    len(self.nodedict[index_to_id[pos]].labels) / self.num_sequences
+                    >= abstention_threshold
+                ):
+                    path.append(index_to_id[pos])
+                    labels.append(self.nodedict[index_to_id[pos]].labels)
+                    text.append(self.nodedict[index_to_id[pos]].text)
+            else:
+                path.append(index_to_id[pos])
+                labels.append(self.nodedict[index_to_id[pos]].labels)
+                text.append(self.nodedict[index_to_id[pos]].text)
+            pos = nextInPath[pos]
+
+        # ignore END node
+        path = path[:-1]
+        # ignore END node
+        text = text[:-1]
+        # ignore START in text
+        text[0] = text[0].replace("START", "")
+        labels = labels[:-1]
+
+        return " ".join(text)
+    
+
+    def consensus_response(
+        self, selection_threshold: Optional[float] = 0.5, api: str = "openai" , model: str = "gpt-4o-mini", task: str = "bio", **kwargs
+    ) -> str:
+        self.toposort()
+
+        consensus_node_ids = self.consensus_node_ids
+        print(consensus_node_ids)
+
+        selected_node_ids = []
+
+        for node_id in consensus_node_ids:
+            if node_id == self.start_id or node_id == self.end_id:
+                continue
+
+            selected_node_ids.append(node_id)
+
+            for neighbor_id in self.nodedict[node_id].outEdges:
+                if neighbor_id in consensus_node_ids:
+                    continue
+
+                if (
+                    len(self.nodedict[neighbor_id].labels) / self.num_sequences
+                    >= selection_threshold
+                ):
+                    selected_node_ids.append(neighbor_id)
+
+        text = " ".join([self.nodedict[node_id].text for node_id in selected_node_ids])
+        print(text)
+        cleaned_text = clean_up_text(text, task=task, api=api, model=model, **kwargs)
+        return cleaned_text
+
+    def save_to_pickle(self, filename):
+        with open(filename, "wb+") as f:
+            pickle.dump(self, f)
+
+    def refine_graph(
+        self,
+        verbose: bool = False,
+        save_intermediate_file: str = None,
+        final_merge: bool = True,
+        **kwargs,
+    ):
+        self.merge_consensus_nodes(verbose=verbose)
+
+        if save_intermediate_file:
+            with open(save_intermediate_file, "w+") as f:
+                self.htmlOutput(f, annotate_consensus=False)
+
+        if not self.consensus_node_ids:
+            self.failed = True
+            return
+
+        else:
+            self.merge_divergent_paths(verbose=verbose, **kwargs)
+
+            if final_merge:
+                try:
+                    self.merge_consensus_nodes(verbose=verbose)
+                except Exception as e:
+                    print(e)
+                    self.failed = True

src/text_poa_graph_utils.py

@@ -0,0 +1,126 @@
+from typing import Optional
+
+from huggingface_hub import InferenceClient
+from openai import OpenAI
+
+
+TEXT_SIMILARITY_JUDGE_PROMPT = """
+You are given two pieces of text. Your task is to determine whether they are semantically equivalent based solely on their factual content.
+
+Here are the specific guidelines:
+- Texts are equivalent if they convey the same core information or concept, regardless of wording or structure
+- If one text has information that is a subset of the other text, then the texts are equivalent
+- Focus ONLY on the essential claims, not on:
+  * Stylistic differences or tone
+  * Level of detail (if the core facts remain the same)
+  * Connotative differences between words
+  * Implied significance or emphasis
+  * Presentation order (if all key information is present in both)
+- Minor additions of non-contradictory information should not make texts non-equivalent
+- For ambiguous cases, prioritize the central claim or purpose of the text
+
+Examples of equivalent pairs:
+- "The meeting starts at 3pm" and "The 3 o'clock meeting will begin on time"
+- "Research indicates a 15% increase" and "Studies show a fifteen percent rise"
+- "was influential in the field" and "had a significant impact on the community"
+
+Examples of non-equivalent pairs:
+- "The project might be completed by Friday" and "The project will be finished by Friday"
+- "Most experts agree on the approach" and "All experts support the approach"
+
+Strictly follow these guidelines and return ONLY:
+- equivalent
+- not equivalent
+"""
+MATH_SIMILARITY_JUDGE_PROMPT = """
+You are given two pieces of text from mathematical solutions. Your task is to determine whether the two solution segments are mathematically equivalent in their content, while allowing for stylistic variations.
+
+Here are some important guidelines:
+- Solutions should be considered equivalent if:
+  1. They communicate the same mathematical content/approach, even if word choice or phrasing differs
+  2. They contain the same key mathematical ideas, even if expressed differently
+  3. The same mathematical steps are described, even if using different words
+  4. They present the same final answer, regardless of wording style or formatting
+
+- Allow for these variations while still considering solutions equivalent:
+  1. Stylistic differences ("we will" vs. "we'll" or "I'll")
+  2. Different levels of formality in the explanation
+  3. Minor rephrasing that preserves the core mathematical content
+  4. Use of synonyms or alternative mathematical terminology for the same concept
+
+- Solutions are NOT equivalent if:
+  1. They use fundamentally different mathematical approaches
+  2. They work with different formulas or equations
+  3. They present different mathematical steps or operations
+  4. They reach different conclusions or answers
+  5. One contains substantial mathematical content that the other lacks
+
+- When examining final answers, focus on mathematical equivalence rather than stylistic presentation
+- For solution steps, maintain the core mathematical approach while allowing for rephrasing
+
+Examples of solutions that SHOULD be considered equivalent:
+- "We will systematically evaluate each possible grouping" and "We'll evaluate each grouping"
+- "The answer is x = 5" and "Therefore, x equals 5"
+- "Using the quadratic formula" and "Applying the quadratic formula"
+
+Strictly follow the guidelines above.
+Return your judgment in the following format. Do not include any other text:
+- equivalent
+- not equivalent
+"""
+
+def path_sim_llm(
+    path1_text: str,
+    path2_text: str,
+    api: str = "openai",
+    model: str = "gpt-4.1-mini",
+    verbose: bool = False,
+    domain: Optional[str] = "text",
+    custom_similarity_judge_prompt: str = None,
+):
+    if api == "openai":
+        client = OpenAI()
+    elif api == "hf":
+        client = InferenceClient()
+    else:
+        raise ValueError(f"Invalid API: {api}")
+
+    if domain == "text":
+        similarity_judge_prompt = (
+            f"{TEXT_SIMILARITY_JUDGE_PROMPT}\n\nText 1: {path1_text}\nText 2: {path2_text}"
+        )
+    elif domain == "math":
+        similarity_judge_prompt = (
+            f"{MATH_SIMILARITY_JUDGE_PROMPT}\n\nText 1: {path1_text}\nText 2: {path2_text}"
+        )
+    elif not domain and custom_similarity_judge_prompt:
+        similarity_judge_prompt = (
+            f"{custom_similarity_judge_prompt}\n\nText 1: {path1_text}\nText 2: {path2_text}"
+        )
+    else:
+        raise ValueError(f"Invalid domain: {domain} and no custom similarity judge prompt provided")
+
+    completion = client.chat.completions.create(
+        model=model,
+        temperature=0,
+        messages=[
+            {"role": "system", "content": "You are a helpful assistant."},
+            {"role": "user", "content": similarity_judge_prompt},
+        ],
+    )
+
+    judgement = completion.choices[0].message.content.strip()
+    judgement = "".join(c for c in judgement if c.isalpha() or c == " ")
+    judgement = judgement.strip()
+
+    if verbose:
+        print(f"{path1_text} \nand \n{path2_text} \nare {judgement}")
+
+    if judgement == "equivalent":
+        return 1, completion.usage.prompt_tokens, completion.usage.completion_tokens
+    elif judgement == "not equivalent":
+        return 0, completion.usage.prompt_tokens, completion.usage.completion_tokens
+    else:
+        if verbose:
+            print(f"Invalid judgement: {judgement}")
+        return 0, completion.usage.prompt_tokens, completion.usage.completion_tokens

src/utils.py

@@ -0,0 +1,46 @@
+from typing import List
+
+from huggingface_hub import InferenceClient
+from openai import OpenAI
+
+
+def detect_abstain(text: str, api: str, model: str):
+    if api == "openai":
+        client = OpenAI()
+    elif api == "hf":
+        client = InferenceClient()
+    else:
+        raise ValueError(f"Invalid API: {api}")
+
+    detect_abstain_prompt = f"""
+    You are given a piece of text that is a part of a biography of an entity. 
+    Text: {text}
+
+    If the text claims a lack of knowledge about the topic, return "Abstain".
+    Otherwise, return "Not abstain".
+    """
+
+    completion = client.chat.completions.create(
+        model=model,
+        messages=[
+            {"role": "system", "content": "You are a helpful assistant."},
+            {"role": "user", "content": detect_abstain_prompt},
+        ],
+    )
+
+    return completion.choices[0].message.content.strip()
+
+
+def calculate_factf1_at_k(
+    supported_facts: List[str], unsupported_facts: List[str], k: int
+) -> float:
+    """
+    Calculate the F1 score at k for supported and unsupported facts
+    """
+    if len(supported_facts) == 0:
+        return 0
+
+    precision = len(supported_facts) / (len(supported_facts) + len(unsupported_facts))
+    recall = min(len(supported_facts) / k, 1)
+    f1 = 2 * precision * recall / (precision + recall)
+    return f1