add paper links

app.py

@@ -13,6 +13,7 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 # TODO Sentence level highlighting instead (prediction after every word is not what it was trained on). Also solves token-level highlighting issues.
 # TODO log prob output scaling highlighting instead?
 # TODO make it look nicer
+# TODO better examples. 
 # TODO streaming output (need custom generation function because of probes)
 # TODO add options to switch between models, SLT/TBG, layers?
 # TODO full semantic entropy calculation
@@ -23,12 +24,18 @@ MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096"))
 
 DESCRIPTION = """
 <h1>Llama-2 7B Chat with Uncertainty Probes</h1>
-<p>This Space demonstrates the Llama-2-7b-chat model with a semantic uncertainty probe.</p>
+<p>This Space demonstrates the Llama-2-7b-chat model with a semantic uncertainty probe. This demo is based on our paper: <a href="https://arxiv.org/abs/2406.15927" target="_blank">"Semantic Uncertainty and Accuracy Probes for Language Models"</a> by Jannik Kossen*, Jiatong Han*, Muhammed Razzak*, Lisa Schut, Shreshth Malik and Yarin Gal.</p>
 <p>The highlighted text shows the model's uncertainty in real-time:</p>
 <ul>
     <li><span style="background-color: #00FF00; color: black">Green</span> indicates more certain generations</li>
     <li><span style="background-color: #FF0000; color: black">Red</span> indicates more uncertain generations</li>
 </ul>
+<p>The demo compares the model's uncertainty with two different probes:</p>
+<ul>
+    <li><b>Semantic Uncertainty Probe:</b> Predicts the semantic uncertainty of the model's generations.</li>
+    <li><b>Accuracy Probe:</b> Predicts the accuracy of the model's generations.</li>
+</ul>
+<p>Please see our paper for more details.</p>
 """
 
 EXAMPLES = [
@@ -101,6 +108,7 @@ def generate(
     se_highlighted_text = ""
     acc_highlighted_text = ""
 
+    # skip the first hidden state as it is the prompt
     for i in range(1, len(hidden)):
 
         # Semantic Uncertainty Probe