import gradio as gr from transformers import AutoModelForSequenceClassification, AutoTokenizer import torch import pickle from dotenv import load_dotenv import os import pandas as pd # Load environment variables from .env file load_dotenv() HF_TOKEN = os.getenv("HF_TOKEN") # Name of dataset to save flagged data to HF_dataset = "peterkros/COFOG-feedback" # <-- Replace with your dataset repo ID # Load the HuggingFaceDatasetSaver logger hf_writer = gr.HuggingFaceDatasetSaver(HF_TOKEN, HF_dataset) level1_to_level2_mapping = { "General public services": [ "Executive and legislative organs, financial and fiscal affairs, external affairs", "Foreign economic aid", "General services", "Basic research", "R&D General public services", "General public services n.e.c.", "Public debt transactions", "Transfers of a general character between different levels of government" ], "Defence": [ "Military defence", "Civil defence", "Foreign military aid", "R&D Defence", "Defence n.e.c." ], "Public order and safety": [ "Police services", "Fire-protection services", "Law courts", "Prisons", "R&D Public order and safety", "Public order and safety n.e.c." ], "Economic affairs": [ "General economic, commercial and labour affairs", "Agriculture, forestry, fishing and hunting", "Fuel and energy", "Mining, manufacturing and construction", "Transport", "Communication", "Other industries", "R&D Economic affairs", "Economic affairs n.e.c." ], "Environmental protection": [ "Waste management", "Waste water management", "Pollution abatement", "Protection of biodiversity and landscape", "R&D Environmental protection", "Environmental protection n.e.c." ], "Housing and community amenities": [ "Housing development", "Community development", "Water supply", "Street lighting", "R&D Housing and community amenities", "Housing and community amenities n.e.c." ], "Health": [ "Medical products, appliances and equipment", "Outpatient services", "Hospital services", "Public health services", "R&D Health", "Health n.e.c." ], "Recreation, culture and religion": [ "Recreational and sporting services", "Cultural services", "Broadcasting and publishing services", "Religious and other community services", "R&D Recreation, culture and religion", "Recreation, culture and religion n.e.c." ], "Education": [ "Pre-primary and primary education", "Secondary education", "Post-secondary non-tertiary education", "Tertiary education", "Education not definable by level", "Subsidiary services to education", "R&D Education", "Education n.e.c." ], "Social protection": [ "Sickness and disability", "Old age", "Survivors", "Family and children", "Unemployment", "Housing", "Social exclusion n.e.c.", "R&D Social protection", "Social protection n.e.c." ] } # Model names for level1 and level2 model_name_level1 = "peterkros/COFOG-bert2" model_name_level2 = "peterkros/COFOG-bert-level2" # Load models and tokenizers for both levels model_level1 = AutoModelForSequenceClassification.from_pretrained(model_name_level1) tokenizer_level1 = AutoTokenizer.from_pretrained(model_name_level1) model_level2 = AutoModelForSequenceClassification.from_pretrained(model_name_level2) tokenizer_level2 = AutoTokenizer.from_pretrained(model_name_level2) # Load the label encoder with open('label_encoder_level1.pkl', 'rb') as file: label_encoder_level1 = pickle.load(file) with open('label_encoder_level2.pkl', 'rb') as file: label_encoder_level2 = pickle.load(file) def predict(text): # Check if the input has at least two words if len(text.split()) < 2: return "Input must have at least two words." # Predict Level1 inputs_level1 = tokenizer_level1(text, return_tensors="pt", padding=True, truncation=True, max_length=512) with torch.no_grad(): outputs_level1 = model_level1(**inputs_level1) probs_level1 = torch.nn.functional.softmax(outputs_level1.logits, dim=-1) predicted_class_level1 = torch.argmax(probs_level1, dim=-1).item() predicted_label_level1 = label_encoder_level1.inverse_transform([predicted_class_level1])[0] # Predict Level2 (assuming level2 model uses both text and predicted level1 label) combined_input = text + " " + predicted_label_level1 inputs_level2 = tokenizer_level2(combined_input, return_tensors="pt", padding=True, truncation=True, max_length=512) with torch.no_grad(): outputs_level2 = model_level2(**inputs_level2) probs_level2 = torch.nn.functional.softmax(outputs_level2.logits, dim=-1) # Extract the probabilities for the candidate level2 categories level2_candidates = level1_to_level2_mapping.get(predicted_label_level1, []) candidate_indices = [label_encoder_level2.transform([candidate])[0] for candidate in level2_candidates if candidate in label_encoder_level2.classes_] # Filter the probabilities filtered_probs = probs_level2[0, candidate_indices] # Get the highest probability label from the filtered list if len(filtered_probs) > 0: highest_prob_index = torch.argmax(filtered_probs).item() predicted_class_level2 = candidate_indices[highest_prob_index] predicted_label_level2 = label_encoder_level2.inverse_transform([predicted_class_level2])[0] else: predicted_label_level2 = "n.e.c" combined_prediction = f"Level1: {predicted_label_level1} - Level2: {predicted_label_level2}" return combined_prediction def classify_csv(file_obj): # Read the CSV file df = pd.read_csv(file_obj) # Check if the 'text' column is in the CSV file if 'text' not in df.columns: return "There is no column named 'text' in the file." # Process the file if the 'text' column exists results = [] for i in range(len(df)): # Combine the current line with the 5 preceding lines for context context_start = max(0, i - 5) context = " ".join(df['text'][context_start:i+1]) # Truncate the context to fit within the model's max length inputs = tokenizer_level1(context, truncation=True, max_length=512, return_tensors="pt") # Extract the truncated text for prediction truncated_context = tokenizer_level1.decode(inputs['input_ids'][0]) # Make a prediction using the truncated context prediction = predict(truncated_context) results.append((df['text'][i], prediction)) # Convert the results to a DataFrame with columns 'Line' and 'Prediction' results_df = pd.DataFrame(results, columns=["Budget Line", "Prediction"]) return results_df # Define the markdown text with bullet points markdown_text = """ - Trained with ~1500 rows of data on bert-base-uncased, English. - Input one budget line per time with min 2 words. - Accuracy of the model is ~88%. """ markdown_text_file_upload = """ - Trained with ~1500 rows of data on bert-base-uncased, English. - Upload CSV ONLY and name your column with budget line item as **text**. - Using RAG (Retrieval-augmented generation) aproach to feed context into classifier using preceding lines of budget. - Accuracy of the model is ~88%. """ html_table = """

COFOG Budget AutoClassification

This classifier was developed utilizing the pre-trained BERT (Bidirectional Encoder Representations from Transformers) model with an uncased configuration, with over 1500 manually labeled dataset comprising budget line items extracted from various budgetary documents. To balance the data, additional data was generated using GPT-4 where categories were not available in budget documents. The model training was executed on a Google Colab environment, specifically utilizing a Tesla T4 GPU. The model is designed to predict the primary classification level of the Classification of the Functions of Government (COFOG), with the predictions from the first level serving as contextual input for subsequent second-level classification. The project is conducted with an exclusive focus on academic and research objectives.
For batch prediction we integrated Retriever-Augmented Generator (RAG) approach. This approach enriches the prediction process by incorporating contextual information from up to 5 preceding lines in the dataset, significantly enhancing the model's ability to understand and classify each entry in the context of related data.
Detailed metrics of the training process are as follows: TrainOutput(global_step=395, training_loss=1.1497593360611156, metrics={'train_runtime': 650.0119, 'train_samples_per_second': 9.638, 'train_steps_per_second': 0.608, 'total_flos': 1648509163714560.0, 'train_loss': 1.1497593360611156, 'epoch': 5.0}).

""" # First interface for single line input iface1 = gr.Interface( fn=predict, inputs=gr.components.Textbox(lines=1, placeholder="Enter Budget line here...", label="Budget Input"), outputs=gr.components.Label(label="Classification Output"), title="COFOG AutoClassification - Single Line", description=markdown_text, article=html_table, allow_flagging="manual", # Enables flagging flagging_options=["Incorect Level1", "Incorect Level2"], flagging_callback=hf_writer, ) # Second interface (for CSV file upload) iface2 = gr.Interface( fn=classify_csv, inputs=gr.components.File(label="Upload CSV File"), outputs=gr.components.DataFrame(label="Classification Results"), description=markdown_text_file_upload, article=html_table, title="COFOG AutoClassification - Batch Classification" ) # Combine the interfaces in a tabbed interface tabbed_interface = gr.TabbedInterface( [iface1, iface2], ["Single Prediction", "Batch Prediction"] ) # Run the interface if __name__ == "__main__": tabbed_interface.launch()