updates

app.py

@@ -22,7 +22,7 @@ print(f"Using device: {device}")
 model_name = "jhu-clsp/FollowIR-7B"
 
 try:
-    model = AutoModelForCausalLM.from_pretrained(model_name).to(device)
+    model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16).to(device)
     tokenizer = AutoTokenizer.from_pretrained(model_name, padding_side="left")
 except ValueError as e:
     print(f"Error loading model or tokenizer: {e}")
@@ -42,6 +42,11 @@ Relevant (only output one word, either "true" or "false"): [/INST] """
 
 @spaces.GPU
 def check_relevance(query, instruction, passage):
+    global model
+    global tokenizer
+    global template
+    global token_false_id
+    global token_true_id
     if torch.cuda.is_available():
         device = "cuda"
         model = model.to(device)
@@ -62,6 +67,7 @@ def check_relevance(query, instruction, passage):
 
     with torch.no_grad():
         batch_scores = model(**tokens).logits[:, -1, :]
+
     true_vector = batch_scores[:, token_true_id]
     false_vector = batch_scores[:, token_false_id]
     batch_scores = torch.stack([false_vector, true_vector], dim=1)
@@ -70,6 +76,25 @@ def check_relevance(query, instruction, passage):
     
     return f"{score:.4f}"
 
+# Example inputs
+examples = [
+    [
+        "What movies were directed by James Cameron?",
+        "A relevant document would describe any movie that was directed by James Cameron",
+        "Avatar: The Way of Water is a 2022 American epic science fiction film co-produced and directed by James Cameron, who co-wrote the screenplay with Rick Jaffa and Amanda Silver."
+    ],
+    [
+        "What are the health benefits of green tea?",
+        "A relevant document would discuss specific health benefits associated with drinking green tea",
+        "Green tea is rich in polyphenols, which are natural compounds that have health benefits, such as reducing inflammation and helping to fight cancer. Green tea contains a catechin called epigallocatechin-3-gallate (EGCG). Catechins are natural antioxidants that help prevent cell damage and provide other benefits."
+    ],
+    [
+        "Who won the Nobel Prize in Physics in 2022?",
+        "A relevant document would mention the names of the physicists who won the Nobel Prize in Physics in 2022",
+        "The 2021 Nobel Prize in Physics was awarded jointly to Syukuro Manabe, Klaus Hasselmann and Giorgio Parisi for groundbreaking contributions to our understanding of complex physical systems."
+    ]
+]
+
 # Gradio Interface
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown("# FollowIR Relevance Checker")
@@ -85,6 +110,14 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
         with gr.Column():
             output = gr.Textbox(label="Relevance Probability")
     
+    gr.Examples(
+        examples=examples,
+        inputs=[query_input, instruction_input, passage_input],
+        outputs=output,
+        fn=check_relevance,
+        cache_examples=True,
+    )
+    
     submit_button.click(
         check_relevance,
         inputs=[query_input, instruction_input, passage_input],