Update app.py

app.py

@@ -1,7 +1,50 @@
 import gradio as gr
+from transformers import AutoModelForSequenceClassification, AutoTokenizer
+from jinja2 import Template
+import torch
 
-def greet(name):
-    return "Hello " + name + "!!"
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()
+# load the judge
+device = "cuda:0"
+model_name = "collinear-ai/collinear-reliability-judge-v1-deberta-ext"
+model = AutoModelForSequenceClassification.from_pretrained(model_name).to(device)
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+
+# tempalte
+template = Template(
+"""
+# Document:
+{{ text }}
+
+# Conversation:
+{{ conversation }}
+"""
+)
+
+
+def judge_reliability(document: str, conversation: str):
+    with torch.no_grad():
+        text = template.render(text=document, conversation=conversation)
+        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[0][1]}"
+    return results
+
+demo = gr.Interface(
+    fn=judge_reliability,
+    inputs=[
+        gr.Textbox(label="Document", lines=5, value="CV was born in Iowa"),
+        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"
+)
+
+if __name__ == "__main__":
+    demo.launch()