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()