app.py

import os
import gradio as gr
from gradio import FlaggingCallback
from gradio.components import IOComponent

from transformers import pipeline

from typing import List, Optional, Any

import argilla as rg

class ArgillaLogger(FlaggingCallback):
    def __init__(self, api_url, api_key, dataset_name):
        rg.init(api_url=api_url, api_key=api_key)
        self.dataset_name = dataset_name
    def setup(self, components: List[IOComponent], flagging_dir: str):
        pass
    def flag(
        self,
        flag_data: List[Any],
        flag_option: Optional[str] = None,
        flag_index: Optional[int] = None,
        username: Optional[str] = None,
    ) -> int:
        text = flag_data[0]
        inference = flag_data[1]
        print(flag_data)
        #rg.log(name=dataset_name,records=rg.TextClassificationRecord(inputs={"answer": text, "response": inference}))



gr.Interface.load(
    "models/deprem-ml/deprem-ner",
    allow_flagging="manual",
    flagging_callback=ArgillaLogger(api_url="https://dvilasuero-argilla-template-space.hf.space", api_key="team.apikey", dataset_name="ner-example")
).launch()