app.py

import gradio as gr
from collinear import Collinear
import os
import json
from openai import AsyncOpenAI
from jinja2 import Template
collinear = Collinear(access_token=os.getenv('COLLINEAR_API_KEY'))
prompt = Template("""
iven the following QUESTION, DOCUMENT and ANSWER you must analyze the provided answer and determine whether it is faithful to the contents of the DOCUMENT. The ANSWER must not offer new information beyond the context provided in the DOCUMENT. The ANSWER also must not contradict information provided in the DOCUMENT. Output your final verdict by strictly following this format: "PASS" if the answer is faithful to the DOCUMENT and "FAIL" if the answer is not faithful to the DOCUMENT. Show your reasoning.
--
QUESTION (THIS DOES NOT COUNT AS BACKGROUND INFORMATION):
{{question}}

--
DOCUMENT:
{{context}}

--
ANSWER:
{{answer}}

--
""")

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)


async def lynx(input_style_dropdown,document_input,question_input,answer_input):
    if input_style_dropdown=='QA format':
        client = AsyncOpenAI(
        base_url="https://s6mipt5j797e6fql.us-east-1.aws.endpoints.huggingface.cloud/v1/", 
        api_key=os.getenv("HF_TOKEN") 
        )
        rendered_prompt = prompt.render(question=question_input,context=document_input,answer=answer_input)
        rendered_prompt +="""
        
Your output should be in JSON FORMAT with the keys "REASONING" and "SCORE":
{{"REASONING": <your reasoning as bullet points>, "SCORE": <your final score>}}
        """
        chat_completion = await client.chat.completions.create(
            model="tgi",
            messages=[
            {
                "role": "user",
                "content": rendered_prompt
            }
        ],
            top_p=None,
            temperature=None,
            max_tokens=150,
            stream=False,
            seed=None,
            frequency_penalty=None,
            presence_penalty=None
        )
        print(chat_completion)
        return chat_completion.choices.pop().message.content
    else:
        return 'NA'

# Function to judge reliability based on the selected input format
async def judge_reliability(input_style, document, conversation, claim, question, answer):
    if input_style == "Conv":
        conversation = json.loads(conversation)
        print(conversation)
        outputs= await collinear.judge.veritas.conversation(document,conversation[:-1],conversation[-1])
    elif input_style == "NLI":
        outputs = await collinear.judge.veritas.natural_language_inference(document,claim)
    elif input_style == "QA format":
        outputs = await collinear.judge.veritas.question_answer(document,question,answer)
    results = f"Reliability Judge Outputs: {outputs}"
    return results



# 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="Alex is a good boy. He stays in California")
        conversation_input = gr.Textbox(label="Conversation", lines=5, visible=True, value='[{"role": "user", "content": "Where is Alex born?"}, {"role": "assistant", "content": "Alex stays in California"}]')
        claim_input = gr.Textbox(label="Claim", lines=5, visible=False, value="Alex stays in California")
        question_input = gr.Textbox(label="Question", lines=5, visible=False, value="Where is Alex born?")
        answer_input = gr.Textbox(label="Answer", lines=5, visible=False, value="Alex stays in California")

    with gr.Row():
        result_output = gr.Textbox(label="Veritas Model")

        lynx_output = gr.Textbox(label="Lynx Model")


    # 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
    ).then(
        fn=lynx,
        inputs=[input_style_dropdown,document_input,question_input,answer_input],
        outputs=lynx_output
    )

# Launch the demo
if __name__ == "__main__":
    demo.launch()