from txtai.pipeline import Labels

class InstructionClassifier:
    def __init__(self):
        # Initialize the labels model
        self.labels = Labels('facebook/bart-large-mnli')

        self.tags = [
            "Programming",
            "Factual",
            "Creative Writing",
            "Roleplaying"
        ]

        self.tools_labels = ["Real Time Information needed: Available in Internet",
                             "Historic Information needed: Available in Wikipedia",
                             "Sufficient Information"]


    def classify_instructions(self, data):
        result = []
        for text in data:
            # Predict tags
            tag_labels_result = self.labels(text, self.tags)
            tag_label = self.tags[tag_labels_result[0][0]] if tag_labels_result[0][0] < len(self.tags) else "Unknown"

            tool_labels_result = self.labels(text, self.tools_labels)
            tool_label = self.tools_labels[tool_labels_result[0][0]] if tool_labels_result[0][0] < len(self.tools_labels) else "Unknown"

            result.append((text, tag_label, tool_label))
        return result
# Usage
classifier = InstructionClassifier()