ICAT-v1-Query / pipelineoperation.py
sankalps's picture
made new text for streamlit spaces
2756ab2
raw
history blame
2.28 kB
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