classify-intent / components /pipeline_preparer.py
aman-s-affinsys's picture
feat: added logging
0c746c9
raw
history blame contribute delete
724 Bytes
from transformers import TextClassificationPipeline
from utils.commons import setup_logging
logger = setup_logging("pipeline_preparer.log")
class PipelinePreparer:
@staticmethod
def prepare_pipeline(tokenizer, model):
"""Create text classification pipeline"""
try:
logger.info("Preparing text classification pipeline...")
pipeline = TextClassificationPipeline(
tokenizer=tokenizer,
model=model,
return_all_scores=True
)
logger.info("Pipeline prepared successfully.")
return pipeline
except Exception as e:
logger.error(f"Error preparing pipeline: {e}")
raise