init

app.py

@@ -0,0 +1,65 @@
+import gradio as gr
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
+
+# Model loading and setup
+model_name = "jhu-clsp/FollowIR-7B"
+model = AutoModelForCausalLM.from_pretrained(model_name).cuda()
+tokenizer = AutoTokenizer.from_pretrained(model_name, padding_side="left")
+tokenizer.pad_token = tokenizer.eos_token
+tokenizer.padding_side = "left"
+token_false_id = tokenizer.get_vocab()["false"]
+token_true_id = tokenizer.get_vocab()["true"]
+
+template = """<s> [INST] You are an expert Google searcher, whose job is to determine if the following document is relevant to the query (true/false). Answer using only one word, one of those two choices.
+
+Query: {query}
+Document: {text}
+Relevant (only output one word, either "true" or "false"): [/INST] """
+
+def check_relevance(query, instruction, passage):
+    full_query = f"{query} {instruction}"
+    prompt = template.format(query=full_query, text=passage)
+    
+    tokens = tokenizer(
+        [prompt],
+        padding=True,
+        truncation=True,
+        return_tensors="pt",
+        pad_to_multiple_of=None,
+    )
+
+    for key in tokens:
+        tokens[key] = tokens[key].cuda()
+
+    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)
+    batch_scores = torch.nn.functional.log_softmax(batch_scores, dim=1)
+    score = batch_scores[:, 1].exp().item()
+    
+    return f"{score:.4f}"
+
+# Gradio Interface
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# FollowIR Relevance Checker")
+    gr.Markdown("This app uses the FollowIR-7B model to determine the relevance of a passage to a given query and instruction.")
+    
+    with gr.Row():
+        with gr.Column():
+            query_input = gr.Textbox(label="Query", placeholder="Enter your search query here")
+            instruction_input = gr.Textbox(label="Instruction", placeholder="Enter additional instructions or criteria")
+            passage_input = gr.Textbox(label="Passage", placeholder="Enter the passage to check for relevance", lines=5)
+            submit_button = gr.Button("Check Relevance")
+        
+        with gr.Column():
+            output = gr.Textbox(label="Relevance Probability")
+    
+    submit_button.click(
+        check_relevance,
+        inputs=[query_input, instruction_input, passage_input],
+        outputs=[output]
+    )
+
+demo.launch()

requirements.txt

@@ -0,0 +1,3 @@
+gradio
+torch
+transformers