from transformers import pipeline
import streamlit as st
pipeline = pipeline(task="text-classification", truncation = True, model="sankalps/NonCompete-Test", token = st.secrets["HF_Token"])
'''
def pipeline_operations(results):
    count_result = -1
    for result in results:
        count_result += 1
        matching_columns = result['matching_columns']
        if len(matching_columns) != 0:
            for matching_column in matching_columns:
                contract_or_not = pipeline(matching_column)
                if contract_or_not == "contractclause":
                    results[count_result]['matching_columns_after_classification'] = matching_column
        elif len(matching_columns) == 0:
            results[count_result]['matching_columns_after_classification'] = ''
    count_result = -1
    for result in results:
        count_result += 1
        matching_indents = result['matching_indents']
        if len(matching_indents) != 0:
            for matching_indent in matching_indents:
                contract_or_not = pipeline(matching_indent)
                if contract_or_not == "contractclause":
                    results[count_result]['matching_indents_after_classification'] = matching_indent
        else:
            results[count_result]['matching_indents_after_classification'] = ''
    return results
'''

def pipeline_operations(results):
    for result in results:
        # Process matching columns
        matching_columns = result.get('matching_columns', [])
        classified_columns = [
            col for col in matching_columns if pipeline(col)[0]['label'] == "contractclause"
        ]
        if classified_columns:
            result['matching_columns_after_classification'] = classified_columns
        else:
            result['matching_columns_after_classification'] = []

        # Process matching indents
        matching_indents = result.get('matching_indents', [])
        classified_indents = [
            indent for indent in matching_indents if pipeline(indent)[0]['label'] == "contractclause"
        ]
        if classified_indents:
            result['matching_indents_after_classification'] = classified_indents
        else:
            result['matching_indents_after_classification'] = []

    return results