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 nlp = pipeline("ner", model="deprem-ml/deprem-ner") examples = [ ["Lütfen yardım Akevler mahallesi Rüzgar sokak Tuncay apartmanı zemin kat Antakya akrabalarım göçük altında #hatay #Afad"] ] def create_record(input_text): # Making the prediction predictions = nlp(input_text, aggregation_strategy="first") # Creating the predicted entities as a list of tuples (entity, start_char, end_char, score) prediction = [(pred["entity_group"], pred["start"], pred["end"], pred["score"]) for pred in predictions] # Create word tokens batch_encoding = nlp.tokenizer(input_text) word_ids = sorted(set(batch_encoding.word_ids()) - {None}) words = [] for word_id in word_ids: char_span = batch_encoding.word_to_chars(word_id) words.append(input_text[char_span.start:char_span.end]) # Building a TokenClassificationRecord record = rg.TokenClassificationRecord( text=input_text, tokens=words, prediction=prediction, prediction_agent="deprem-ml/deprem-ner", ) print(record) return record 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] rg.log(name=self.dataset_name, records=create_record(text)) gr.Interface.load( "models/deprem-ml/deprem-ner", examples=examples, allow_flagging="manual", flagging_callback=ArgillaLogger( api_url="https://dvilasuero-argilla-template-1-3.hf.space", api_key="team.apikey", dataset_name="ner-gradio-feedback" ), #flaggging_options=["Correct", "Incorrect", "Ambiguous"] ).launch()