app.py

import gradio as gr
from transformers import AutoModelForSequenceClassification, AutoTokenizer
from jinja2 import Template
import torch


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