from obsei.payload import TextPayload
from transformers import pipeline
from obsei.analyzer.classification_analyzer import ZeroShotClassificationAnalyzer, ClassificationAnalyzerConfig

# Step 1: Define your configuration (labels, etc.)
analyzer_config = ClassificationAnalyzerConfig(
    labels=["Sports", "Politics", "Technology", "Entertainment"],  # Example labels
    multi_class_classification=False,
    add_positive_negative_labels=False
)

# Step 2: Initialize the ZeroShotClassificationAnalyzer
analyzer = ZeroShotClassificationAnalyzer(
    model_name_or_path="facebook/bart-large-mnli",  # Using a pre-trained zero-shot classification model
    device="cpu"  # Use GPU, set to "cpu" if using CPU
)

# Step 3: Prepare the input text (as TextPayload objects)
texts = [
    "The new iPhone has been released and it's taking the tech world by storm.",
    "The latest political debate had strong views on the economy.",
    "The football match between Barcelona and Madrid ended in a draw."
]

# Create TextPayloads from the texts
source_responses = [TextPayload(processed_text=text) for text in texts]

# Step 4: Run the analysis
results = analyzer.analyze_input(source_response_list=source_responses, analyzer_config=analyzer_config)

# Step 5: Output the results
for result in results:
    print(f"Text: {result.processed_text}")
    print(f"Classification Scores: {result.segmented_data['classifier_data']}")