Update app.py

app.py

@@ -4,52 +4,112 @@ from jinja2 import Template
 import torch
 import json
 
-
 # load the judge
 device = "cuda:0"
 model_name = "collinear-ai/collinear-reliability-judge-v5"
 model = AutoModelForSequenceClassification.from_pretrained(model_name).to(device)
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 
+# templates
+conv_template = Template(
+    """
+    # Document:
+    {{ document }}
+
+    # Conversation:
+    {% for message in conversation %}
+    {{ message.role }}: {{ message.content }}
+    {% endfor %}
+    """
+)
 
-# template
-template = Template(
-"""
-# Document:
-{{ document }}
+qa_template = Template(
+    """
+    # Document:
+    {{ document }}
 
-# Conversation:
-{% for message in conversation %}
-{{ message.role }}: {{ message.content }}
-{% endfor %}
-"""
+    # Question:
+    {{ question }}
+
+    # Answer:
+    {{ answer }}
+    """
 )
 
+nli_template = Template(
+    """
+    # Document:
+    {{ document }}
+
+    # Claim:
+    {{ claim }}
+    """
+)
+
+
+# Function to dynamically update inputs based on the input style
+def update_inputs(input_style):
+    if input_style == "Conv":
+        return gr.update(visible=True), gr.update(visible=True), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
+    elif input_style == "NLI":
+        return gr.update(visible=True), gr.update(visible=False), gr.update(visible=True), gr.update(visible=False), gr.update(visible=False)
+    elif input_style == "QA format":
+        return gr.update(visible=True), gr.update(visible=False), gr.update(visible=False), gr.update(visible=True), gr.update(visible=True)
+
 
-def judge_reliability(document: str, conversation: str):
+# Function to judge reliability based on the selected input format
+def judge_reliability(input_style, document, conversation, claim, question, answer):
     with torch.no_grad():
-        conversation = json.loads(conversation)
-        text = template.render(document=document, conversation=conversation)
+        if input_style == "Conv":
+            conversation = json.loads(conversation)
+            text = conv_template.render(document=document, conversation=conversation)
+        elif input_style == "NLI":
+            text = nli_template.render(document=document, claim=claim)
+        elif input_style == "QA format":
+            text = qa_template.render(document=document, question=question, answer=answer)
+
         print(text)
+        
         encoded = tokenizer([text], padding=True)
         input_ids = torch.tensor(encoded.input_ids).to(device)
         attention_mask = torch.tensor(encoded.attention_mask).to(device)
         outputs = model.forward(input_ids=input_ids, attention_mask=attention_mask)
         outputs = torch.softmax(outputs.logits, axis=1)
-        results = f"Reliability Score: {outputs}"
+        results = f"Reliability Score: {outputs[0][1].item()}"
     return results
 
-demo = gr.Interface(
-    fn=judge_reliability,
-    inputs=[
-        gr.Textbox(label="Document", lines=5, value="Chris Voss, was born in Iowa, USA. He is the best negotiator in the world."),
-        gr.Textbox(label="Conversation", lines=5, value='[{"role": "user", "content": "Where are you born?"}, {"role": "assistant", "content": "I am born in Iowa"}]')
-    ],
-    outputs=gr.Textbox(label="Results"),
-    title="Collinear Reliability Judge",
-    description="Enter a document and conversation (json formatted) to judge reliability. Note: this judges if the last assistant turn is faithful according to the given document ",
-    theme="default"
-)
 
+
+# Create the interface using gr.Blocks
+with gr.Blocks() as demo:
+    with gr.Row():
+        input_style_dropdown = gr.Dropdown(label="Input Style", choices=["Conv", "NLI", "QA format"], value="Conv", visible=True)
+
+    with gr.Row():
+        document_input = gr.Textbox(label="Document", lines=5, visible=True, value="Chris Voss is one of the best negotiators in the world. And he was born in Iowa, USA.")
+        conversation_input = gr.Textbox(label="Conversation", lines=5, visible=True, value='[{"role": "user", "content": "Hi Chris Voss, Where are you born?"}, {"role": "assistant", "content": "I am born in Iowa"}]')
+        claim_input = gr.Textbox(label="Claim", lines=5, visible=False, value="CV was born in Iowa")
+        question_input = gr.Textbox(label="Question", lines=5, visible=False, value="Where is Chris Voss born?")
+        answer_input = gr.Textbox(label="Answer", lines=5, visible=False, value="CV was born in Iowa")
+
+    with gr.Row():
+        result_output = gr.Textbox(label="Results")
+
+
+    # Set the visibility of inputs based on the selected input style
+    input_style_dropdown.change(
+        fn=update_inputs, 
+        inputs=[input_style_dropdown], 
+        outputs=[document_input, conversation_input, claim_input, question_input, answer_input]
+    )
+
+    # Set the function to handle the reliability check
+    gr.Button("Submit").click(
+        fn=judge_reliability, 
+        inputs=[input_style_dropdown, document_input, conversation_input, claim_input, question_input, answer_input], 
+        outputs=result_output
+    )
+
+# Launch the demo
 if __name__ == "__main__":
     demo.launch()