Organisation

app.py

@@ -1,88 +1,13 @@
 import gradio as gr
-import math
-import os
-import numpy as np
-from huggingface_hub import InferenceClient
 
-client = InferenceClient(token=os.environ.get("HF_TOKEN"))
+from src.rendering import (
+    LEGEND_HTML,
+    TOOLTIP_CSS
+)
 
-def hex_to_rgb(hex_code):
-    r = int(hex_code[0:2], 16)
-    g = int(hex_code[2:4], 16)
-    b = int(hex_code[4:6], 16)
-    return r, g, b
+from src.generation import query
 
-def prob_to_color(avg_prob):
-    norm_prob = np.clip(avg_prob, 0.0, 1.0)
-    r = int((1 - norm_prob) * 255)
-    g = int(norm_prob * 255)
-    b = 0
-    return f"{r:02x}{g:02x}{b:02x}"
 
-def tokens_to_html(logprobs_content):
-    spans = ""
-    for token in logprobs_content:
-        prob = math.exp(token.logprob)
-        hex_color = prob_to_color(prob)
-        r, g, b = hex_to_rgb(hex_color)
-        brightness = (r * 299 + g * 587 + b * 114) / 1000
-        text_color = "black" if brightness > 128 else "white"
-        token_text = token.token.replace(" ", "&nbsp;").replace("<", "&lt;").replace(">", "&gt;")
-        tooltip_text = f"p={prob:.3f} logprob={token.logprob:.3f}"
-        spans += (
-            f'<span '
-            f'data-tooltip="{tooltip_text}" '
-            f'style="background-color:#{hex_color}; color:{text_color}; '
-            f'padding:2px 4px; border-radius:3px; margin:1px; cursor:pointer;">'
-            f'{token_text}</span>'
-        )
-    return f"<p style='font-family: monospace; font-size: 1.1em; line-height: 2em;'>{spans}</p>"
-
-LEGEND_HTML = """
-<div style='margin:10px 0; font-family:monospace;'>
-  <b>Token confidence:</b>&nbsp;
-  <span style='background:#ff0000; color:white; padding:4px 8px; border-radius:3px;'>Low (0%)</span>&nbsp;
-  <span style='background:#7f7f00; color:white; padding:4px 8px; border-radius:3px;'>Moderate (50%)</span>&nbsp;
-  <span style='background:#00ff00; color:black; padding:4px 8px; border-radius:3px;'>High (100%)</span>
-</div>
-"""
-
-TOOLTIP_CSS = """
-span[data-tooltip] {
-    position: relative;
-}
-span[data-tooltip]:hover::after {
-    content: attr(data-tooltip);
-    position: absolute;
-    bottom: 130%;
-    left: 50%;
-    transform: translateX(-50%);
-    background: #333;
-    color: white;
-    padding: 4px 8px;
-    border-radius: 4px;
-    font-size: 0.75em;
-    white-space: nowrap;
-    z-index: 9999;
-    pointer-events: none;
-}
-"""
-
-def query(user_input, model, max_tokens):
-    if not user_input.strip():
-        return "<p>Please enter a question.</p>"
-    try:
-        response = client.chat.completions.create(
-            model=model,
-            messages=[{"role": "user", "content": user_input}],
-            max_tokens=int(max_tokens),
-            logprobs=True,
-            top_logprobs=1
-        )
-        print(f"QUERY: {user_input}")
-        return tokens_to_html(response.choices[0].logprobs.content)
-    except Exception as e:
-        return f"<p style='color:red;'>Error: {str(e)}</p>"
 
 with gr.Blocks(title="Probabot", css=TOOLTIP_CSS) as demo:
     gr.Markdown("# 🤖🔍 Probabot")

utils.py → src/entailment.py

@@ -1,18 +1,9 @@
 import os
-import logging
 import numpy as np
 import torch
 import torch.nn.functional as F
 from transformers import AutoModelForSequenceClassification, AutoTokenizer
 
-import math
-
-import json
-import pandas as pd
-import matplotlib.pyplot as plt
-from matplotlib.cm import get_cmap
-import random
-
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
 
@@ -97,87 +88,4 @@ def get_semantic_ids(strings_list, model, strict_entailment=False, example=None)
     return semantic_ids
     
 
-'''
-LLM CHECKS ENTAILMENT
-'''
-def LLM_get_semantic_ids(strings_list, sub_question, MODEL_NAME, strict_entailment=False):
-    def LLM_check_implication(text1, text2):
-        '''
-        LLM checks if text1 implies text2. Returns 2 if entailment, 1 if neutral and 0 if contradiction.
-        '''
-        try:
-            client = genai.Client()
-        except Exception as e:
-            print(f"Error when initializing client: {e}")
-            if not os.environ.get("GEMINI_API_KEY"):
-                print("ERROR: GEMINI_API_KEY not found in environment.")
-            return
-
-        # the prompt below was copied from semantic_uncertainty/semantic_uncertainty/uncertainty/uncertainty_measures/semantic_entropy.py
-        prompt = f"""We are evaluating answers to the question \"{sub_question}\"\n"""
-        prompt += "Here are two possible answers:\n"
-        prompt += f"Possible Answer 1: {text1}\nPossible Answer 2: {text2}\n"
-        prompt += "Does Possible Answer 1 semantically entail Possible Answer 2? Respond with entailment, contradiction, or neutral, nothing else. """
-
-        try:
-            response = client.models.generate_content(
-                model=MODEL_NAME,  
-                contents=prompt, 
-                config=GenerateContentConfig(
-                    temperature = 0.7 # temperature = 0 -> models behave deterministically. temperature > 0 -> model randomly choses from top tokens
-                )
-            )
-            entailment_answer = response.lower()[:30]
-            if 'entailment' in entailment_answer:
-                return 2 
-            elif 'neutral' in entailment_answer:
-                return 1
-            elif 'contradiction' in entailment_answer:
-                return 0
-            else:
-                raise ValueError(
-                    f"Unexpected model output: {entailment_answer}. Answer can only be 'entailment, 'neutral' or 'contradiction'. "
-                )
-        
-        except Exception as e:
-            raise RuntimeError(
-                f"API call failed for prompt: {prompt}"
-            ) from e
-
-    def LLM_are_equivalent(text1, text2):
-        '''
-        Returns True if text1 and text2 are semantically identical. 
-        Two strings text1 and text2 are semantically equivalent (with strict_entailment=False) when 
-        their implications are either [1, 2], [2, 1] or [2, 2]. With strict_entailment=True, then only [2, 2] means they are equivalent
-        '''
-        try: 
-            i1 = LLM_check_implication(text1, text2)
-        except ValueError as e:
-            print(f"LLM failed in checking implication of \n'{text1}' --> '{text2}'. \nThe error that was raised is: \n{e}")
-
-        try:
-            i2 = LLM_check_implication(text2, text1)
-        except ValueError as e:
-            print(f"LLM failed in checking implication of \n'{text2}' --> '{text1}'. \nThe error that was raised is: \n{e}")
-
-        if strict_entailment:
-            return (i1 == 2) and (i2 == 2)
-        else:
-            implications = [i1, i2]
-            return (0 not in implications) and (implications != [1, 1])
-
-    semantic_ids = [-1] * len(strings_list)
-    next_id = 0
-
-    # give each string in strings_list a semnatic_id
-    for i, s1 in enumerate(strings_list):
-        if semantic_ids[i] == -1:
-            semantic_ids[i] = next_id
-            for j in range(i + 1, len(strings_list)):
-                if LLM_are_equivalent(s1, strings_list[j]):
-                    semantic_ids[j] = next_id # give the string the same semantic id if it is equivalent with an other string in strings_list
-            next_id += 1
-
-    return semantic_ids
-
     

src/generation.py

@@ -0,0 +1,330 @@
+import os 
+import requests
+import re
+import json
+import time
+
+def query_model(system_prompt, prompt, model_name, max_tokens, client=None):
+    """
+    Parameters
+    ----------
+    system_prompt : str   — the system prompt
+    prompt        : str   — the input text
+    model_name    : str   — "gpt2", "llama-1b", or "llama-8b"
+    max_tokens    : int   — how many tokens to generate (capped at 200 by proxy)
+    client        : str   — the proxy URL returned by setup_client()
+
+    Returns
+    -------
+    dict with keys:
+        "answer"      : str   — the generated text
+        "logprobs"    : dict  — {token: logprob}
+        "token_probs" : dict  — {token: probability}
+    """
+    if client is None:
+        client = "https://cryotud-nb4170-llm-proxy.hf.space" 
+
+    # Also include SYSTEM PROMPT!
+    response = requests.post(
+        f"{client}/generate",
+        json={
+            "system_prompt": system_prompt, 
+            "prompt": prompt,
+            "model": model_name,
+            "max_tokens": max_tokens,
+        },
+        timeout=60,
+    )
+
+    response.raise_for_status()
+    return response.json()
+
+## HIGH LEVEL FUNCTIONS ##
+def generate_long_answer(question, model_name, max_tokens, client):
+    '''
+    Generates a longer answer of three sentences to the provided main question. 
+    Returns a string that contains the answer.
+    '''
+
+    prompt = question + '\nAnswer in exactly 3 sentences. Do not use bullet points in your answer.'
+    response = query_model('', prompt, model_name, max_tokens, client)
+    answer = response['answer']
+    return answer
+
+## LEVEL BELOW ## 
+
+def split_answers_into_json_strings(answer, model_name, max_tokens, client, verbose=True):
+    """
+    Makes 1 API call to split the answer into sentences, then 1 API call per sentence 
+    to extract sentence parts and factoids
+    Returns raw parsed JSON
+    """
+
+    sys_prompt_1 = """"You are a strict text parser. 
+    You ONLY output valid JSON arrays of strings. 
+    No markdown, no explanations.
+    Example Input: 
+    Text: Paris is the capital of France. The Eiffel Tower is in Paris. The Louvre is a famous museum in Paris.
+    Example Output:
+    ["Paris is the capital of France.", "The Eiffel Tower is in Paris.", "The Louvre is a famous museum in Paris."]
+    """
+
+    user_prompt_1 = f"""Split the following text into individual sentences.
+    Rules:
+    1. Output a maximum of 3 sentences.
+    2. Respond ONLY with a valid JSON array of strings.
+    
+    Text: "{answer}"
+    
+    Example Output:
+    ["Sentence 1.", "Sentence 2.", "Sentence 3."]
+    """
+
+
+    if verbose:
+        print("\n" + "=" * 80)
+        print("SENTENCE SPLITTING (JSON)")
+        print("=" * 80)
+
+    # First API call: split the sentences using LLM
+    try:
+        res_sentences = query_model(sys_prompt_1, user_prompt_1, model_name, max_tokens, client)
+        raw_sentences = res_sentences.get('answer', '').strip()
+        clean_json_str = re.sub(r'```json|```', '', raw_sentences).strip() # safety cleaner to process potential "```json" in answer
+        sentences = json.loads(clean_json_str)
+        print(f"Successfully extracted {len(sentences)} sentences via JSON.")
+    except Exception as e:
+        print(f"❌ Error during sentence splitting: {e}")
+        return [], []
+    
+    #####  FACTOID DECOMPOSITION #####
+    if verbose:
+        print("\n" + "=" * 80)
+        print("FACTOID EXTRACTION (JSON)")
+        print("=" * 80)
+        
+    master_json_output = []
+
+    sys_prompt_2 = """You are a strict data extraction tool. 
+    You ONLY output valid JSON arrays of objects.
+    You NEVER invent external information.
+    
+    EXAMPLE INPUT:
+    Sentence: "Marie Curie was born in Warsaw, and she later moved to Paris to study."
+    
+    EXAMPLE OUTPUT:
+    [
+        {{
+        "sentence_part": "Marie Curie was born in Warsaw, ",
+        "factoids": ["Marie Curie was born in Warsaw"]
+        }},
+        {{
+        "sentence_part": "and she later moved to Paris to study.",
+        "factoids": ["Marie Curie moved to Paris", "Marie Curie moved to Paris for her studies"]
+        }}
+    ]
+    """	
+    
+    # API calls 2-4: split into 2 sentence parts, and split those into 2 factoids
+    for i, sentence in enumerate(sentences, 1):
+        print(f"Processing Sentence {i}/{len(sentences)}: '{sentence}'")
+        
+        user_prompt_2 = f"""Decompose the given sentence into consecutive 'sentence parts' and extract factoids.
+                
+        RULES:
+        1. Min 1, max 2 sentence parts total.
+        2. EXACT COPY: The 'sentence part' MUST be an exact, unaltered substring copied directly from the original sentence.
+        3. Min 1, max 2 brief factoids per sentence part.
+        4. STRICT: ONLY extract facts explicitly stated.
+        5. Respond ONLY with a valid JSON array.
+        6. Be exact in the factoids that you generate: provide names of the people/places/occurences that you mention.
+        
+        EXAMPLE INPUT:
+        Sentence: "Marie Curie was born in Warsaw, and she later moved to Paris to study."
+        
+        EXAMPLE OUTPUT:
+        [
+          {{
+            "sentence_part": "Marie Curie was born in Warsaw, ",
+            "factoids": ["Marie Curie was born in Warsaw"]
+          }},
+          {{
+            "sentence_part": "and she later moved to Paris to study.",
+            "factoids": ["Marie Curie moved to Paris", "Marie Curie moved to Paris for her studies"]
+          }}
+        ]
+        
+        ACTUAL INPUT:
+        Sentence: "{sentence}"
+        """
+        
+        try:
+            res_factoids = query_model(sys_prompt_2, user_prompt_2, model_name, max_tokens, client)
+            raw_output = res_factoids.get('answer', '').strip()
+            clean_output = re.sub(r'```json|```', '', raw_output).strip() # again safety check to remove additional json stuff
+            
+            # Parse the JSON and append to master list
+            parsed_sentence_data = json.loads(clean_output) # json.loads() parses the json string and puts the contents into a dictionary
+            master_json_output.append(parsed_sentence_data) 
+            
+        except Exception as e:
+            print(f"❌ Error parsing JSON for sentence {i}: {e}")
+            master_json_output.append([]) # Append empty list on failure to keep index aligned
+            
+    return master_json_output, sentences
+
+def parse_json_to_dictionaries(master_json_output, original_sentences):
+    """
+    Takes the raw JSON output from the LLM pipeline, applies Python code with safety limits,
+    and puts it into the designated dictionaries
+    """
+    sentence_part_to_factoids = {}
+    factoid_to_sentence_part = {}
+    sentence_parts = []
+    factoids_list = []
+    
+    for sentence_data, original_sentence in zip(master_json_output, original_sentences):
+        constrained_parts = sentence_data[:2] # ensure there are only two sentence parts
+        
+        for item in constrained_parts:
+            s_part = item.get("sentence_part", "").strip() # get the sentence part
+            extracted_factoids = item.get("factoids", []) # get the corresponding factoids
+            
+            if not s_part:
+                continue
+                
+            if s_part not in original_sentence:
+                print(f"⚠️ WARNING: Paraphrase detected! \nExpected substring of: '{original_sentence}' \nGot: '{s_part}'")
+            
+            constrained_factoids = [f.strip() for f in extracted_factoids if f.strip()][:2] # strip of whitespace & ensure there are only two factoids
+
+            # print warning if a sentence part has no extracted factiods
+            if len(constrained_factoids) == 0:
+                print(f"⚠️ Notice: No valid factoids extracted for sentence part: '{s_part}'")
+                continue
+
+            # store the sentence part and the factoids in the corresponding datastructures
+            sentence_parts.append(s_part) 
+            sentence_part_to_factoids[s_part] = constrained_factoids
+            
+            for factoid in constrained_factoids:
+                factoids_list.append(factoid)
+                factoid_to_sentence_part[factoid] = s_part
+
+    return sentence_part_to_factoids, factoid_to_sentence_part, sentence_parts, factoids_list
+
+
+## LOW LEVEL FUNCTIONS ##
+def generate_Q_questions(factoid, original_answer, Q, model_name, max_tokens, client, enable_sleep=True):
+    '''
+    Generates Q questions that might have generated the factoid, using the original answer for context.
+    
+    Args:
+        - factoid: string, the current factoid to turn into questions
+        - original_answer: string, the original full text
+        - Q: int, number of questions to generate
+        - model_name: string, name of the model
+        - max_tokens: int
+        - client: the proxy URL
+        - enable_sleep: boolean, enables sleep function if rate limit is hit
+        
+    Returns:
+        - questions: a list of exactly Q question strings
+    '''
+
+
+    sys_prompt = """You are a precise data generator. You ONLY output valid JSON arrays of strings. No explanations, no markdown, no conversational text."""
+    
+    user_prompt = f"""Based on the Original Text, generate exactly {Q} different questions that are directly answered by the provided factoid.
+    
+    Original Text: "{original_answer}"
+    
+    Factoid to target: "{factoid}"
+    
+    Rules:
+    1. Generate exactly {Q} questions.
+    2. Vary the phrasing of the questions.
+    3. Be specific: ALWAYS include the names of the objects, people, or places in your questions so they can be understood without seeing the original text.
+    4. Do not ask questions about information that is NOT in the factoid.
+    5. Respond ONLY with a valid JSON array of strings. 
+
+    Example Input:
+    "Marie Curie moved to Paris for her studies"
+    
+    Example Output:
+    ["Why did Marie Curie move to Paris?", "Where did Marie Curie move to for her studies?"]
+    """
+    
+    # Retry loop: query the model and if the rate limit is hit, wait for 20 secs with a maximum of 3 retries
+    max_retries = 3
+    for attempt in range(max_retries):
+        try:
+            response = query_model(sys_prompt, user_prompt, model_name, max_tokens, client)
+            raw_output = response.get('answer', '').strip() # extract and strip the answer
+            clean_json_str = re.sub(r'```json|```', '', raw_output).strip() # remove potential json mark
+            questions = json.loads(clean_json_str) # load the answer into a dictionary
+            constrained_questions = [str(q).strip() for q in questions][:Q] # strip and ensure that there are max Q questions
+            return constrained_questions
+            
+        except Exception as e:
+            error_str = str(e).lower()
+
+            # if the rate limit is hit, sleep for 20 secs
+            if enable_sleep and ('429' in error_str or 'rate limit' in error_str or 'too many requests' in error_str):
+                wait_time = 20 # Wait 20 seconds before trying again
+                print(f"⏳ Rate limit hit! Sleeping for {wait_time}s... (Attempt {attempt+1}/{max_retries})")
+                time.sleep(wait_time)
+            else:
+                print(f"❌ ERROR: API Error for factoid '{factoid[:30]}...': {e}")
+                if attempt == max_retries - 1 or not enable_sleep:
+                    return []
+    return []
+
+def generate_M_answers(question, M, model_name, max_tokens, client, enable_sleep=True):
+    '''
+    Provides M independent answers to the current question to measure model uncertainty.
+    
+    Args:
+        - question: string, the question to answer
+        - M: int, number of times to ask the question
+        - model_name: string, name of the model
+        - max_tokens: int
+        - client: the proxy URL
+        - enable_sleep: boolean, enables sleep function if rate limit is hit
+        
+    Returns:
+        - sub_answers: list of M answer strings
+    '''
+    
+    sys_prompt = "You are a highly concise question-answering system. You ONLY output the direct answer entity. No full sentences, no conversational text."
+    user_prompt = f"""Question: {question}
+    
+    Rules:
+    1. Do NOT answer in a full sentence.
+    2. Answer with as few words as possible (e.g., only a name, place, number, or thing).
+    3. If multiple names/places/things are applicable, provide a EXACTLY 1 item.
+    4. Provide ONLY the answer and absolutely nothing else.
+    """
+    sub_answers = []
+    
+    for i in range(M):
+        # Again enable 3 retries if rate limit is hit 
+        max_retries = 3
+        for attempt in range(max_retries):
+            try:
+                response = query_model(sys_prompt, user_prompt, model_name, max_tokens, client)
+                raw_answer = response.get('answer', '').strip()
+                sub_answers.append(raw_answer)
+                break # If a response is generated, break out of the retry loop and go to the next 'M'
+            except Exception as e:
+                error_str = str(e).lower()
+                if enable_sleep and ('429' in error_str or 'rate limit' in error_str or 'too many requests' in error_str):
+                    wait_time = 20
+                    print(f"⏳ Rate limit hit! Sleeping for {wait_time}s... (Attempt {attempt+1}/{max_retries})")
+                    time.sleep(wait_time)
+                else:
+                    print(f"❌ Error during API call {i+1} for question: {e}")
+                    if attempt == max_retries - 1 or not enable_sleep:
+                        sub_answers.append("ERROR") 
+                        break
+    return sub_answers

src/rendering.py

@@ -0,0 +1,65 @@
+
+
+LEGEND_HTML = """
+<div style='margin:10px 0; font-family:monospace;'>
+  <b>Token confidence:</b>&nbsp;
+  <span style='background:#ff0000; color:white; padding:4px 8px; border-radius:3px;'>Low (0%)</span>&nbsp;
+  <span style='background:#7f7f00; color:white; padding:4px 8px; border-radius:3px;'>Moderate (50%)</span>&nbsp;
+  <span style='background:#00ff00; color:black; padding:4px 8px; border-radius:3px;'>High (100%)</span>
+</div>
+"""
+
+TOOLTIP_CSS = """
+span[data-tooltip] {
+    position: relative;
+}
+span[data-tooltip]:hover::after {
+    content: attr(data-tooltip);
+    position: absolute;
+    bottom: 130%;
+    left: 50%;
+    transform: translateX(-50%);
+    background: #333;
+    color: white;
+    padding: 4px 8px;
+    border-radius: 4px;
+    font-size: 0.75em;
+    white-space: nowrap;
+    z-index: 9999;
+    pointer-events: none;
+}
+"""
+
+def hex_to_rgb(hex_code):
+    r = int(hex_code[0:2], 16)
+    g = int(hex_code[2:4], 16)
+    b = int(hex_code[4:6], 16)
+    return r, g, b
+
+def prob_to_color(avg_prob):
+    import numpy as np
+    norm_prob = np.clip(avg_prob, 0.0, 1.0)
+    r = int((1 - norm_prob) * 255)
+    g = int(norm_prob * 255)
+    b = 0
+    return f"{r:02x}{g:02x}{b:02x}"
+
+def tokens_to_html(logprobs_content):
+    import math
+    spans = ""
+    for token in logprobs_content:
+        prob = math.exp(token.logprob)
+        hex_color = prob_to_color(prob)
+        r, g, b = hex_to_rgb(hex_color)
+        brightness = (r * 299 + g * 587 + b * 114) / 1000
+        text_color = "black" if brightness > 128 else "white"
+        token_text = token.token.replace(" ", "&nbsp;").replace("<", "&lt;").replace(">", "&gt;")
+        tooltip_text = f"p={prob:.3f} logprob={token.logprob:.3f}"
+        spans += (
+            f'<span '
+            f'data-tooltip="{tooltip_text}" '
+            f'style="background-color:#{hex_color}; color:{text_color}; '
+            f'padding:2px 4px; border-radius:3px; margin:1px; cursor:pointer;">'
+            f'{token_text}</span>'
+        )
+    return f"<p style='font-family: monospace; font-size: 1.1em; line-height: 2em;'>{spans}</p>"