ollama-classifier 0.1.0__py3-none-any.whl

ollama_classifier/__init__.py

@@ -0,0 +1,30 @@
+"""Ollama Classifier - A wrapper around Ollama Python SDK for text classification.
+
+This package provides a classifier for text classification with constrained output
+and confidence scoring using the Ollama Python SDK.
+
+Example usage:
+    from ollama import Client
+    from ollama_classifier import OllamaClassifier, ClassificationResult
+    
+    client = Client()
+    classifier = OllamaClassifier(client, model="llama3.2")
+    
+    result = classifier.classify(
+        text="I love this product!",
+        choices=["positive", "negative", "neutral"]
+    )
+    print(f"Prediction: {result.prediction}")
+    print(f"Confidence: {result.confidence:.2%}")
+"""
+
+from .types import ClassificationResult, ChoicesType
+from .classifier import OllamaClassifier
+
+__all__ = [
+    "OllamaClassifier",
+    "ClassificationResult",
+    "ChoicesType",
+]
+
+__version__ = "0.1.0"

ollama_classifier/classifier.py

@@ -0,0 +1,845 @@
+"""Main classifier module for ollama-classifier."""
+
+import json
+import math
+from typing import Union, List, Dict, Any
+
+from ollama import Client, AsyncClient
+
+from .types import ClassificationResult, ChoicesType
+from .prompts import (
+    build_classification_prompt,
+    get_choice_labels,
+    build_json_schema_for_choices,
+)
+
+
+class OllamaClassifier:
+    """A classifier wrapper around Ollama client for text classification.
+    
+    This class provides methods for classifying text into a set of predefined choices
+    with support for both fast (single-call) and complete (multi-call) scoring methods.
+    
+    Attributes:
+        _client: The Ollama client (sync or async).
+        _model: The model name to use for classification.
+    """
+    
+    def __init__(self, client: Union[Client, AsyncClient], model: str):
+        """Initialize the classifier.
+        
+        Args:
+            client: An initialized Ollama Client or AsyncClient instance.
+            model: The model name to use for classification (e.g., "llama3.2").
+        """
+        self._client = client
+        self._model = model
+    
+    # =========================================================================
+    # Sync Methods - Generate
+    # =========================================================================
+    
+    def generate(
+        self,
+        text: str,
+        choices: ChoicesType,
+        system_prompt: str | None = None,
+    ) -> str:
+        """Generate a constrained classification for a single text.
+        
+        Uses JSON schema with enum constraint to ensure only valid choices are generated.
+        This is the fastest method as it only makes one API call and doesn't compute
+        confidence scores.
+        
+        Args:
+            text: The text to classify.
+            choices: Either a list of choice labels, or a dict mapping labels to descriptions.
+            system_prompt: Optional custom system prompt.
+            
+        Returns:
+            The predicted choice label.
+        """
+        labels = get_choice_labels(choices)
+        system, user = build_classification_prompt(text, choices, system_prompt)
+        schema = build_json_schema_for_choices(labels)
+        
+        messages = [
+            {"role": "system", "content": system},
+            {"role": "user", "content": user},
+        ]
+        
+        response = self._client.chat(
+            model=self._model,
+            messages=messages,
+            format=schema,
+            options={"temperature": 0.0},
+        )
+        
+        content = response.get("message", {}).get("content", "{}")
+        result = json.loads(content)
+        return result.get("label", "")
+    
+    def batch_generate(
+        self,
+        texts: List[str],
+        choices: ChoicesType,
+        system_prompt: str | None = None,
+    ) -> List[str]:
+        """Generate constrained classifications for multiple texts.
+        
+        Args:
+            texts: List of texts to classify.
+            choices: Either a list of choice labels, or a dict mapping labels to descriptions.
+            system_prompt: Optional custom system prompt.
+            
+        Returns:
+            List of predicted choice labels, one per input text.
+        """
+        return [self.generate(text, choices, system_prompt) for text in texts]
+    
+    # =========================================================================
+    # Sync Methods - Score Fast (Single-call logprob extraction)
+    # =========================================================================
+    
+    def score_fast(
+        self,
+        text: str,
+        choices: ChoicesType,
+        system_prompt: str | None = None,
+    ) -> ClassificationResult:
+        """Score a classification using single-call logprob extraction.
+        
+        Makes one API call with logprobs enabled and extracts the probability
+        distribution from the token distribution at the decision point.
+        This is faster but may be less reliable for capturing all choice probabilities.
+        
+        Args:
+            text: The text to classify.
+            choices: Either a list of choice labels, or a dict mapping labels to descriptions.
+            system_prompt: Optional custom system prompt.
+            
+        Returns:
+            ClassificationResult with prediction, confidence, and probabilities.
+        """
+        labels = get_choice_labels(choices)
+        system, user = build_classification_prompt(text, choices, system_prompt)
+        schema = build_json_schema_for_choices(labels)
+        
+        messages = [
+            {"role": "system", "content": system},
+            {"role": "user", "content": user},
+        ]
+        
+        response = self._client.chat(
+            model=self._model,
+            messages=messages,
+            format=schema,
+            logprobs=True,
+            top_logprobs=20,
+            options={"temperature": 0.0},
+        )
+        
+        # Extract prediction from response
+        content = response.get("message", {}).get("content", "{}")
+        result = json.loads(content)
+        prediction = result.get("label", "")
+        
+        # Extract logprobs from token distribution
+        probabilities = self._extract_probabilities_from_logprobs(
+            response, labels
+        )
+        
+        return ClassificationResult(
+            prediction=prediction,
+            confidence=probabilities.get(prediction, 0.0),
+            probabilities=probabilities,
+            raw_response=response,
+        )
+    
+    def batch_score_fast(
+        self,
+        texts: List[str],
+        choices: ChoicesType,
+        system_prompt: str | None = None,
+    ) -> List[ClassificationResult]:
+        """Score multiple texts using fast single-call method.
+        
+        Args:
+            texts: List of texts to classify.
+            choices: Either a list of choice labels, or a dict mapping labels to descriptions.
+            system_prompt: Optional custom system prompt.
+            
+        Returns:
+            List of ClassificationResults, one per input text.
+        """
+        return [self.score_fast(text, choices, system_prompt) for text in texts]
+    
+    # =========================================================================
+    # Sync Methods - Score Complete (Multi-call evaluation)
+    # =========================================================================
+    
+    def score_complete(
+        self,
+        text: str,
+        choices: ChoicesType,
+        system_prompt: str | None = None,
+    ) -> ClassificationResult:
+        """Score a classification using multi-call evaluation with softmax.
+        
+        Makes separate API calls for each choice to compute log P(choice|context),
+        then applies softmax for calibrated probabilities.
+        This is more accurate but slower (N API calls for N choices).
+        
+        Args:
+            text: The text to classify.
+            choices: Either a list of choice labels, or a dict mapping labels to descriptions.
+            system_prompt: Optional custom system prompt.
+            
+        Returns:
+            ClassificationResult with prediction, confidence, and probabilities.
+        """
+        labels = get_choice_labels(choices)
+        system, user = build_classification_prompt(text, choices, system_prompt)
+        
+        # Compute logprob for each choice
+        logprobs: Dict[str, float] = {}
+        for label in labels:
+            logprobs[label] = self._get_choice_logprob(
+                system, user, label
+            )
+        
+        # Apply softmax to get probabilities
+        probabilities = self._softmax(logprobs)
+        
+        # Get prediction (highest probability)
+        prediction = max(probabilities, key=probabilities.get)
+        
+        return ClassificationResult(
+            prediction=prediction,
+            confidence=probabilities[prediction],
+            probabilities=probabilities,
+            raw_response={"logprobs": logprobs},
+        )
+    
+    def batch_score_complete(
+        self,
+        texts: List[str],
+        choices: ChoicesType,
+        system_prompt: str | None = None,
+    ) -> List[ClassificationResult]:
+        """Score multiple texts using complete multi-call method.
+        
+        Args:
+            texts: List of texts to classify.
+            choices: Either a list of choice labels, or a dict mapping labels to descriptions.
+            system_prompt: Optional custom system prompt.
+            
+        Returns:
+            List of ClassificationResults, one per input text.
+        """
+        return [self.score_complete(text, choices, system_prompt) for text in texts]
+    
+    # =========================================================================
+    # Sync Methods - Classify (Generate + Score)
+    # =========================================================================
+    
+    def classify(
+        self,
+        text: str,
+        choices: ChoicesType,
+        system_prompt: str | None = None,
+    ) -> ClassificationResult:
+        """Classify text with constrained generation and fast scoring.
+        
+        Combines generate() and score_fast() for a single classification
+        with confidence scores.
+        
+        Args:
+            text: The text to classify.
+            choices: Either a list of choice labels, or a dict mapping labels to descriptions.
+            system_prompt: Optional custom system prompt.
+            
+        Returns:
+            ClassificationResult with prediction, confidence, and probabilities.
+        """
+        return self.score_fast(text, choices, system_prompt)
+    
+    def classify_complete(
+        self,
+        text: str,
+        choices: ChoicesType,
+        system_prompt: str | None = None,
+    ) -> ClassificationResult:
+        """Classify text with constrained generation and complete scoring.
+        
+        Combines generate() and score_complete() for accurate classification
+        with calibrated confidence scores.
+        
+        Args:
+            text: The text to classify.
+            choices: Either a list of choice labels, or a dict mapping labels to descriptions.
+            system_prompt: Optional custom system prompt.
+            
+        Returns:
+            ClassificationResult with prediction, confidence, and probabilities.
+        """
+        # Use generate for constrained prediction
+        prediction = self.generate(text, choices, system_prompt)
+        
+        # Use score_complete for probabilities
+        result = self.score_complete(text, choices, system_prompt)
+        
+        # Override prediction with constrained output
+        result.prediction = prediction
+        result.confidence = result.probabilities.get(prediction, 0.0)
+        
+        return result
+    
+    def batch_classify(
+        self,
+        texts: List[str],
+        choices: ChoicesType,
+        system_prompt: str | None = None,
+    ) -> List[ClassificationResult]:
+        """Classify multiple texts with fast scoring.
+        
+        Args:
+            texts: List of texts to classify.
+            choices: Either a list of choice labels, or a dict mapping labels to descriptions.
+            system_prompt: Optional custom system prompt.
+            
+        Returns:
+            List of ClassificationResults, one per input text.
+        """
+        return [self.classify(text, choices, system_prompt) for text in texts]
+    
+    def batch_classify_complete(
+        self,
+        texts: List[str],
+        choices: ChoicesType,
+        system_prompt: str | None = None,
+    ) -> List[ClassificationResult]:
+        """Classify multiple texts with complete scoring.
+        
+        Args:
+            texts: List of texts to classify.
+            choices: Either a list of choice labels, or a dict mapping labels to descriptions.
+            system_prompt: Optional custom system prompt.
+            
+        Returns:
+            List of ClassificationResults, one per input text.
+        """
+        return [self.classify_complete(text, choices, system_prompt) for text in texts]
+    
+    # =========================================================================
+    # Async Methods - Generate
+    # =========================================================================
+    
+    async def agenerate(
+        self,
+        text: str,
+        choices: ChoicesType,
+        system_prompt: str | None = None,
+    ) -> str:
+        """Async version of generate().
+        
+        Generate a constrained classification for a single text.
+        
+        Args:
+            text: The text to classify.
+            choices: Either a list of choice labels, or a dict mapping labels to descriptions.
+            system_prompt: Optional custom system prompt.
+            
+        Returns:
+            The predicted choice label.
+        """
+        labels = get_choice_labels(choices)
+        system, user = build_classification_prompt(text, choices, system_prompt)
+        schema = build_json_schema_for_choices(labels)
+        
+        messages = [
+            {"role": "system", "content": system},
+            {"role": "user", "content": user},
+        ]
+        
+        response = await self._client.chat(
+            model=self._model,
+            messages=messages,
+            format=schema,
+            options={"temperature": 0.0},
+        )
+        
+        content = response.get("message", {}).get("content", "{}")
+        result = json.loads(content)
+        return result.get("label", "")
+    
+    async def abatch_generate(
+        self,
+        texts: List[str],
+        choices: ChoicesType,
+        system_prompt: str | None = None,
+    ) -> List[str]:
+        """Async version of batch_generate().
+        
+        Generate constrained classifications for multiple texts.
+        
+        Args:
+            texts: List of texts to classify.
+            choices: Either a list of choice labels, or a dict mapping labels to descriptions.
+            system_prompt: Optional custom system prompt.
+            
+        Returns:
+            List of predicted choice labels, one per input text.
+        """
+        import asyncio
+        return await asyncio.gather(*[
+            self.agenerate(text, choices, system_prompt) for text in texts
+        ])
+    
+    # =========================================================================
+    # Async Methods - Score Fast
+    # =========================================================================
+    
+    async def ascore_fast(
+        self,
+        text: str,
+        choices: ChoicesType,
+        system_prompt: str | None = None,
+    ) -> ClassificationResult:
+        """Async version of score_fast().
+        
+        Score a classification using single-call logprob extraction.
+        
+        Args:
+            text: The text to classify.
+            choices: Either a list of choice labels, or a dict mapping labels to descriptions.
+            system_prompt: Optional custom system prompt.
+            
+        Returns:
+            ClassificationResult with prediction, confidence, and probabilities.
+        """
+        labels = get_choice_labels(choices)
+        system, user = build_classification_prompt(text, choices, system_prompt)
+        schema = build_json_schema_for_choices(labels)
+        
+        messages = [
+            {"role": "system", "content": system},
+            {"role": "user", "content": user},
+        ]
+        
+        response = await self._client.chat(
+            model=self._model,
+            messages=messages,
+            format=schema,
+            logprobs=True,
+            top_logprobs=20,
+            options={"temperature": 0.0},
+        )
+        
+        # Extract prediction from response
+        content = response.get("message", {}).get("content", "{}")
+        result = json.loads(content)
+        prediction = result.get("label", "")
+        
+        # Extract logprobs from token distribution
+        probabilities = self._extract_probabilities_from_logprobs(
+            response, labels
+        )
+        
+        return ClassificationResult(
+            prediction=prediction,
+            confidence=probabilities.get(prediction, 0.0),
+            probabilities=probabilities,
+            raw_response=response,
+        )
+    
+    async def abatch_score_fast(
+        self,
+        texts: List[str],
+        choices: ChoicesType,
+        system_prompt: str | None = None,
+    ) -> List[ClassificationResult]:
+        """Async version of batch_score_fast().
+        
+        Score multiple texts using fast single-call method.
+        
+        Args:
+            texts: List of texts to classify.
+            choices: Either a list of choice labels, or a dict mapping labels to descriptions.
+            system_prompt: Optional custom system prompt.
+            
+        Returns:
+            List of ClassificationResults, one per input text.
+        """
+        import asyncio
+        return await asyncio.gather(*[
+            self.ascore_fast(text, choices, system_prompt) for text in texts
+        ])
+    
+    # =========================================================================
+    # Async Methods - Score Complete
+    # =========================================================================
+    
+    async def ascore_complete(
+        self,
+        text: str,
+        choices: ChoicesType,
+        system_prompt: str | None = None,
+    ) -> ClassificationResult:
+        """Async version of score_complete().
+        
+        Score a classification using multi-call evaluation with softmax.
+        
+        Args:
+            text: The text to classify.
+            choices: Either a list of choice labels, or a dict mapping labels to descriptions.
+            system_prompt: Optional custom system prompt.
+            
+        Returns:
+            ClassificationResult with prediction, confidence, and probabilities.
+        """
+        import asyncio
+        
+        labels = get_choice_labels(choices)
+        system, user = build_classification_prompt(text, choices, system_prompt)
+        
+        # Compute logprob for each choice concurrently
+        logprobs_tasks = [
+            self._aget_choice_logprob(system, user, label)
+            for label in labels
+        ]
+        logprob_values = await asyncio.gather(*logprobs_tasks)
+        
+        logprobs = dict(zip(labels, logprob_values))
+        
+        # Apply softmax to get probabilities
+        probabilities = self._softmax(logprobs)
+        
+        # Get prediction (highest probability)
+        prediction = max(probabilities, key=probabilities.get)
+        
+        return ClassificationResult(
+            prediction=prediction,
+            confidence=probabilities[prediction],
+            probabilities=probabilities,
+            raw_response={"logprobs": logprobs},
+        )
+    
+    async def abatch_score_complete(
+        self,
+        texts: List[str],
+        choices: ChoicesType,
+        system_prompt: str | None = None,
+    ) -> List[ClassificationResult]:
+        """Async version of batch_score_complete().
+        
+        Score multiple texts using complete multi-call method.
+        
+        Args:
+            texts: List of texts to classify.
+            choices: Either a list of choice labels, or a dict mapping labels to descriptions.
+            system_prompt: Optional custom system prompt.
+            
+        Returns:
+            List of ClassificationResults, one per input text.
+        """
+        import asyncio
+        return await asyncio.gather(*[
+            self.ascore_complete(text, choices, system_prompt) for text in texts
+        ])
+    
+    # =========================================================================
+    # Async Methods - Classify
+    # =========================================================================
+    
+    async def aclassify(
+        self,
+        text: str,
+        choices: ChoicesType,
+        system_prompt: str | None = None,
+    ) -> ClassificationResult:
+        """Async version of classify().
+        
+        Classify text with constrained generation and fast scoring.
+        
+        Args:
+            text: The text to classify.
+            choices: Either a list of choice labels, or a dict mapping labels to descriptions.
+            system_prompt: Optional custom system prompt.
+            
+        Returns:
+            ClassificationResult with prediction, confidence, and probabilities.
+        """
+        return await self.ascore_fast(text, choices, system_prompt)
+    
+    async def aclassify_complete(
+        self,
+        text: str,
+        choices: ChoicesType,
+        system_prompt: str | None = None,
+    ) -> ClassificationResult:
+        """Async version of classify_complete().
+        
+        Classify text with constrained generation and complete scoring.
+        
+        Args:
+            text: The text to classify.
+            choices: Either a list of choice labels, or a dict mapping labels to descriptions.
+            system_prompt: Optional custom system prompt.
+            
+        Returns:
+            ClassificationResult with prediction, confidence, and probabilities.
+        """
+        # Use generate for constrained prediction
+        prediction = await self.agenerate(text, choices, system_prompt)
+        
+        # Use score_complete for probabilities
+        result = await self.ascore_complete(text, choices, system_prompt)
+        
+        # Override prediction with constrained output
+        result.prediction = prediction
+        result.confidence = result.probabilities.get(prediction, 0.0)
+        
+        return result
+    
+    async def abatch_classify(
+        self,
+        texts: List[str],
+        choices: ChoicesType,
+        system_prompt: str | None = None,
+    ) -> List[ClassificationResult]:
+        """Async version of batch_classify().
+        
+        Classify multiple texts with fast scoring.
+        
+        Args:
+            texts: List of texts to classify.
+            choices: Either a list of choice labels, or a dict mapping labels to descriptions.
+            system_prompt: Optional custom system prompt.
+            
+        Returns:
+            List of ClassificationResults, one per input text.
+        """
+        import asyncio
+        return await asyncio.gather(*[
+            self.aclassify(text, choices, system_prompt) for text in texts
+        ])
+    
+    async def abatch_classify_complete(
+        self,
+        texts: List[str],
+        choices: ChoicesType,
+        system_prompt: str | None = None,
+    ) -> List[ClassificationResult]:
+        """Async version of batch_classify_complete().
+        
+        Classify multiple texts with complete scoring.
+        
+        Args:
+            texts: List of texts to classify.
+            choices: Either a list of choice labels, or a dict mapping labels to descriptions.
+            system_prompt: Optional custom system prompt.
+            
+        Returns:
+            List of ClassificationResults, one per input text.
+        """
+        import asyncio
+        return await asyncio.gather(*[
+            self.aclassify_complete(text, choices, system_prompt) for text in texts
+        ])
+    
+    # =========================================================================
+    # Private Helper Methods
+    # =========================================================================
+    
+    def _extract_probabilities_from_logprobs(
+        self,
+        response: Dict[str, Any],
+        labels: List[str],
+    ) -> Dict[str, float]:
+        """Extract probabilities from logprobs in the response.
+        
+        Looks for the decision token in the logprobs and extracts
+        the probability distribution for the valid choices.
+        
+        Args:
+            response: The raw API response.
+            labels: List of valid choice labels.
+            
+        Returns:
+            Dict mapping choice labels to probabilities.
+        """
+        # Initialize with zero logprobs
+        choice_logprobs: Dict[str, float] = {label: float("-inf") for label in labels}
+        
+        # Navigate the logprobs payload to find the decision token
+        logprobs_data = response.get("logprobs", [])
+        
+        if not logprobs_data:
+            # No logprobs available, return uniform distribution
+            return {label: 1.0 / len(labels) for label in labels}
+        
+        for token_data in logprobs_data:
+            # Check if this is a dict or list format
+            if isinstance(token_data, dict):
+                top_candidates = token_data.get("top_logprobs", [])
+            elif isinstance(token_data, (list, tuple)) and len(token_data) > 0:
+                # Might be in different format
+                top_candidates = token_data if isinstance(token_data[0], dict) else []
+            else:
+                continue
+            
+            found_decision_token = False
+            
+            for candidate in top_candidates:
+                if isinstance(candidate, dict):
+                    token = candidate.get("token", "")
+                    # Normalize the token
+                    clean_token = token.strip(' \n\r\t"\'').lower()
+                    
+                    # Check if this matches any of our labels
+                    for label in labels:
+                        if label.lower() == clean_token or clean_token in label.lower():
+                            logprob = candidate.get("logprob", float("-inf"))
+                            if logprob > choice_logprobs[label]:
+                                choice_logprobs[label] = logprob
+                            found_decision_token = True
+            
+            if found_decision_token:
+                break
+        
+        # Apply softmax to convert logprobs to probabilities
+        return self._softmax(choice_logprobs)
+    
+    def _get_choice_logprob(
+        self,
+        system: str,
+        user: str,
+        choice: str,
+    ) -> float:
+        """Compute log P(choice | context) for a single choice.
+        
+        Appends the choice as assistant response and measures
+        how "expected" those tokens were.
+        
+        Args:
+            system: System prompt.
+            user: User prompt.
+            choice: The choice label to evaluate.
+            
+        Returns:
+            Log probability of the choice given the context.
+        """
+        messages = [
+            {"role": "system", "content": system},
+            {"role": "user", "content": user},
+            {"role": "assistant", "content": json.dumps({"label": choice})},
+        ]
+        
+        response = self._client.chat(
+            model=self._model,
+            messages=messages,
+            logprobs=True,
+            options={"num_predict": 0},
+        )
+        
+        return self._extract_logprob_from_response(response)
+    
+    async def _aget_choice_logprob(
+        self,
+        system: str,
+        user: str,
+        choice: str,
+    ) -> float:
+        """Async version of _get_choice_logprob().
+        
+        Compute log P(choice | context) for a single choice.
+        
+        Args:
+            system: System prompt.
+            user: User prompt.
+            choice: The choice label to evaluate.
+            
+        Returns:
+            Log probability of the choice given the context.
+        """
+        messages = [
+            {"role": "system", "content": system},
+            {"role": "user", "content": user},
+            {"role": "assistant", "content": json.dumps({"label": choice})},
+        ]
+        
+        response = await self._client.chat(
+            model=self._model,
+            messages=messages,
+            logprobs=True,
+            options={"num_predict": 0},
+        )
+        
+        return self._extract_logprob_from_response(response)
+    
+    def _extract_logprob_from_response(self, response: Dict[str, Any]) -> float:
+        """Extract log probability from Ollama response.
+        
+        Tries various fields where Ollama might return logprobs.
+        
+        Args:
+            response: The raw API response.
+            
+        Returns:
+            Extracted log probability, or 0.0 if not found.
+        """
+        msg = response.get("message", {})
+        
+        # Field: message.logprob (most common in newer versions)
+        if "logprob" in msg:
+            return msg["logprob"]
+        
+        # Field: eval_prob
+        if "eval_prob" in response:
+            p = response["eval_prob"]
+            return math.log(p) if p > 0 else float("-inf")
+        
+        # Field: logprobs list
+        if "logprobs" in msg:
+            lp = msg["logprobs"]
+            if isinstance(lp, list) and lp:
+                # Sum or average the logprobs
+                return sum(lp) / len(lp)
+        
+        # Check in response logprobs
+        logprobs = response.get("logprobs", [])
+        if logprobs and isinstance(logprobs, list):
+            # Try to extract from first token's logprob
+            if isinstance(logprobs[0], dict) and "logprob" in logprobs[0]:
+                return logprobs[0]["logprob"]
+        
+        return 0.0  # Default neutral
+    
+    def _softmax(self, logprobs: Dict[str, float]) -> Dict[str, float]:
+        """Apply numerically stable softmax to log probabilities.
+        
+        Args:
+            logprobs: Dict mapping labels to log probabilities.
+            
+        Returns:
+            Dict mapping labels to probabilities (summing to 1.0).
+        """
+        # Filter out -inf values for computation
+        valid_logprobs = {k: v for k, v in logprobs.items() if v > float("-inf")}
+        
+        if not valid_logprobs:
+            # If all are -inf, return uniform distribution
+            n = len(logprobs)
+            return {k: 1.0 / n for k in logprobs}
+        
+        max_lp = max(valid_logprobs.values())
+        exp_vals = {k: math.exp(v - max_lp) if v > float("-inf") else 0.0 
+                    for k, v in logprobs.items()}
+        total = sum(exp_vals.values())
+        
+        if total == 0:
+            n = len(logprobs)
+            return {k: 1.0 / n for k in logprobs}
+        
+        return {k: v / total for k, v in exp_vals.items()}

ollama_classifier/prompts.py

@@ -0,0 +1,97 @@
+"""Prompt building utilities for classification."""
+
+from typing import List, Dict, Union
+
+
+def build_classification_prompt(
+    text: str,
+    choices: Union[List[str], Dict[str, str]],
+    system_prompt: str | None = None,
+) -> tuple[str, str]:
+    """Build the system and user prompts for classification.
+    
+    Args:
+        text: The text to classify.
+        choices: Either a list of choice labels, or a dict mapping labels to descriptions.
+        system_prompt: Optional custom system prompt.
+        
+    Returns:
+        Tuple of (system_prompt, user_prompt) strings.
+    """
+    # Build the choices section of the prompt
+    choices_text = _format_choices(choices)
+    
+    # Default system prompt
+    if system_prompt is None:
+        system_prompt = (
+            "You are a precise text classifier. "
+            "Your task is to classify the given text into exactly one of the provided categories. "
+            "Respond with only the category label, nothing else."
+        )
+    
+    # Build user prompt
+    user_prompt = f"""Classify the following text into one of these categories:
+
+{choices_text}
+
+Text to classify:
+{text}
+
+Respond with only the category label."""
+    
+    return system_prompt, user_prompt
+
+
+def _format_choices(choices: Union[List[str], Dict[str, str]]) -> str:
+    """Format choices for the prompt.
+    
+    Args:
+        choices: Either a list of choice labels, or a dict mapping labels to descriptions.
+        
+    Returns:
+        Formatted string representation of choices.
+    """
+    if isinstance(choices, dict):
+        # Choices with descriptions
+        lines = []
+        for label, description in choices.items():
+            lines.append(f"- {label}: {description}")
+        return "\n".join(lines)
+    else:
+        # Simple list of choices
+        return "\n".join(f"- {choice}" for choice in choices)
+
+
+def get_choice_labels(choices: Union[List[str], Dict[str, str]]) -> List[str]:
+    """Extract the choice labels from either format.
+    
+    Args:
+        choices: Either a list of choice labels, or a dict mapping labels to descriptions.
+        
+    Returns:
+        List of choice labels.
+    """
+    if isinstance(choices, dict):
+        return list(choices.keys())
+    return list(choices)
+
+
+def build_json_schema_for_choices(choices: List[str]) -> dict:
+    """Build a JSON schema that constrains output to the given choices.
+    
+    Args:
+        choices: List of valid choice labels.
+        
+    Returns:
+        JSON schema dict for structured output.
+    """
+    return {
+        "type": "object",
+        "properties": {
+            "label": {
+                "type": "string",
+                "enum": choices,
+            }
+        },
+        "required": ["label"],
+    }

ollama_classifier/types.py

@@ -0,0 +1,24 @@
+"""Type definitions for ollama-classifier."""
+
+from dataclasses import dataclass, field
+from typing import Dict, Any, Union, List
+
+
+@dataclass
+class ClassificationResult:
+    """Result of a classification operation.
+    
+    Attributes:
+        prediction: The predicted class label.
+        confidence: Confidence score for the prediction (0.0 to 1.0).
+        probabilities: Probability distribution over all choices.
+        raw_response: Raw response from Ollama API for debugging.
+    """
+    prediction: str
+    confidence: float
+    probabilities: Dict[str, float]
+    raw_response: Dict[str, Any] = field(default_factory=dict)
+
+
+# Type alias for choices - can be a list of labels or dict mapping labels to descriptions
+ChoicesType = Union[List[str], Dict[str, str]]

ollama_classifier-0.1.0.dist-info/METADATA

@@ -0,0 +1,247 @@
+Metadata-Version: 2.3
+Name: ollama-classifier
+Version: 0.1.0
+Summary: A wrapper around Ollama Python SDK for text classification with constrained output and confidence scoring
+Author: Luigi Palumbo, Mengting Yu
+Author-email: Luigi Palumbo <paluigi@gmail.com>, Mengting Yu <mengting.yu@unitus.it>
+Requires-Dist: ollama>=0.4.0
+Requires-Dist: sphinx>=7.0.0 ; extra == 'docs'
+Requires-Dist: sphinx-rtd-theme>=2.0.0 ; extra == 'docs'
+Requires-Dist: myst-parser>=2.0.0 ; extra == 'docs'
+Requires-Python: >=3.11
+Provides-Extra: docs
+Description-Content-Type: text/markdown
+
+# ollama-classifier
+
+A Python wrapper around the Ollama Python SDK for text classification with constrained output and confidence scoring.
+
+## Features
+
+- **Constrained Output**: Uses JSON schema with enum constraints to ensure only valid choices are generated
+- **Confidence Scoring**: Two methods available:
+  - **Fast**: Single API call with logprob extraction
+  - **Complete**: Multi-call evaluation with softmax for calibrated probabilities
+- **Sync & Async**: Full support for both synchronous and asynchronous operations
+- **Batch Processing**: Classify multiple texts efficiently
+- **Flexible Choices**: Support for simple labels or labels with descriptions
+- **Custom Prompts**: Override the default system prompt for specialized tasks
+
+## Installation
+
+```bash
+pip install ollama-classifier
+```
+
+Or with uv:
+
+```bash
+uv add ollama-classifier
+```
+
+## Prerequisites
+
+- [Ollama](https://ollama.com/download) installed and running
+- A model pulled (e.g., `ollama pull llama3.2`)
+
+## Quick Start
+
+```python
+from ollama import Client
+from ollama_classifier import OllamaClassifier
+
+client = Client()
+classifier = OllamaClassifier(client, model="llama3.2")
+
+result = classifier.classify(
+    text="I love this product!",
+    choices=["positive", "negative", "neutral"]
+)
+
+print(f"Prediction: {result.prediction}")
+print(f"Confidence: {result.confidence:.2%}")
+print(f"Probabilities: {result.probabilities}")
+```
+
+## Usage
+
+### Basic Classification
+
+```python
+from ollama import Client
+from ollama_classifier import OllamaClassifier
+
+client = Client()
+classifier = OllamaClassifier(client, model="llama3.2")
+
+result = classifier.classify(
+    text="The goalkeeper made an incredible save!",
+    choices=["sports", "politics", "technology", "entertainment"]
+)
+```
+
+### Classification with Label Descriptions
+
+Providing descriptions helps the model understand each category better:
+
+```python
+choices = {
+    "positive": "Text expresses happiness, satisfaction, or approval",
+    "negative": "Text expresses anger, disappointment, or disapproval",
+    "mixed": "Text contains both positive and negative sentiments",
+    "neutral": "Text is factual without strong emotional content",
+}
+
+result = classifier.classify(
+    text="The food was amazing but the service was terrible.",
+    choices=choices
+)
+```
+
+### Custom System Prompt
+
+```python
+result = classifier.classify(
+    text="The quarterly earnings exceeded analyst expectations.",
+    choices=["bullish", "bearish", "neutral"],
+    system_prompt="You are a financial sentiment analyzer. "
+                  "Classify financial news based on market sentiment."
+)
+```
+
+### Scoring Methods
+
+#### Fast Scoring (Single API Call)
+
+```python
+result = classifier.score_fast(
+    text="The movie was fantastic!",
+    choices=["positive", "negative", "neutral"]
+)
+```
+
+#### Complete Scoring (Multi-Call with Softmax)
+
+More accurate but makes N API calls for N choices:
+
+```python
+result = classifier.score_complete(
+    text="The movie was fantastic!",
+    choices=["positive", "negative", "neutral"]
+)
+```
+
+### Generate Only (Fastest)
+
+When you only need the prediction without confidence scores:
+
+```python
+prediction = classifier.generate(
+    text="The team won the championship!",
+    choices=["sports", "finance", "politics"]
+)
+```
+
+### Batch Classification
+
+```python
+texts = [
+    "The goalkeeper made an incredible save!",
+    "The central bank raised interest rates.",
+    "The new smartphone features a revolutionary camera.",
+]
+
+# Fast batch classification
+results = classifier.batch_classify(
+    texts=texts,
+    choices=["sports", "finance", "technology"]
+)
+
+# Complete batch classification (more accurate)
+results = classifier.batch_classify_complete(
+    texts=texts,
+    choices=["sports", "finance", "technology"]
+)
+
+for text, result in zip(texts, results):
+    print(f"{text} -> {result.prediction} ({result.confidence:.2%})")
+```
+
+### Async Usage
+
+```python
+import asyncio
+from ollama import AsyncClient
+from ollama_classifier import OllamaClassifier
+
+async def main():
+    client = AsyncClient()
+    classifier = OllamaClassifier(client, model="llama3.2")
+    
+    # Single classification
+    result = await classifier.aclassify(
+        text="The concert was amazing!",
+        choices=["positive", "negative", "neutral"]
+    )
+    
+    # Batch classification (concurrent)
+    results = await classifier.abatch_classify(
+        texts=["Text 1", "Text 2", "Text 3"],
+        choices=["positive", "negative", "neutral"]
+    )
+
+asyncio.run(main())
+```
+
+## API Reference
+
+### ClassificationResult
+
+```python
+@dataclass
+class ClassificationResult:
+    prediction: str              # The predicted choice label
+    confidence: float            # Confidence score (0.0 to 1.0)
+    probabilities: Dict[str, float]  # Probability distribution over all choices
+    raw_response: Dict           # Raw Ollama response for debugging
+```
+
+### OllamaClassifier Methods
+
+| Method | Async | Description |
+|--------|-------|-------------|
+| `generate(text, choices, system_prompt)` | `agenerate` | Constrained output only (fastest) |
+| `score_fast(text, choices, system_prompt)` | `ascore_fast` | Single-call logprob extraction |
+| `score_complete(text, choices, system_prompt)` | `ascore_complete` | Multi-call evaluation with softmax |
+| `classify(text, choices, system_prompt)` | `aclassify` | Generate + score_fast |
+| `classify_complete(text, choices, system_prompt)` | `aclassify_complete` | Generate + score_complete |
+| `batch_generate(texts, choices, system_prompt)` | `abatch_generate` | Batch constrained output |
+| `batch_score_fast(texts, choices, system_prompt)` | `abatch_score_fast` | Batch fast scoring |
+| `batch_score_complete(texts, choices, system_prompt)` | `abatch_score_complete` | Batch complete scoring |
+| `batch_classify(texts, choices, system_prompt)` | `abatch_classify` | Batch classify (fast) |
+| `batch_classify_complete(texts, choices, system_prompt)` | `abatch_classify_complete` | Batch classify (complete) |
+
+### Parameters
+
+- **text** (str): The text to classify
+- **texts** (List[str]): List of texts to classify (batch methods)
+- **choices** (Union[List[str], Dict[str, str]]): Either a list of choice labels, or a dict mapping labels to descriptions
+- **system_prompt** (str | None): Optional custom system prompt
+
+## Choosing a Method
+
+| Use Case | Recommended Method |
+|----------|-------------------|
+| Speed is critical, no confidence needed | `generate` |
+| Speed with confidence scores | `classify` / `score_fast` |
+| Accurate confidence scores | `classify_complete` / `score_complete` |
+| Batch processing | `batch_classify` or `batch_classify_complete` |
+| Concurrent processing | Async variants (`aclassify`, etc.) |
+
+## License
+
+MIT License
+
+## Development
+
+This project just started! Looking forward to suggestions, issues, and pull requests!

ollama_classifier-0.1.0.dist-info/RECORD

@@ -0,0 +1,8 @@
+ollama_classifier/__init__.py,sha256=r5LFD-YlcVQKMJ1RUTtaZV6WlyHCFdsWwFN-12tKhLE,862
+ollama_classifier/classifier.py,sha256=4XwDRqJ8EIbnNlL8Zp0TP0Ff6QMRJkYczxhiDe2jL40,30140
+ollama_classifier/prompts.py,sha256=54YuHN7941LXtY-S7-krfPurzzvULunCbA84PMyVhk0,2746
+ollama_classifier/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+ollama_classifier/types.py,sha256=UU9R8TksYfjKKJ1ZRoyr5FTsFlCv2rLx2Kt6K2-P1o0,773
+ollama_classifier-0.1.0.dist-info/WHEEL,sha256=mydTeHxOpFHo-DnYhAd_3ATePms-g4rrYvM7wJK8P-U,80
+ollama_classifier-0.1.0.dist-info/METADATA,sha256=JEHvA1nLU5JYE0g1v35GfqnTt6_dqjKjQiQHCRq_YYw,7414
+ollama_classifier-0.1.0.dist-info/RECORD,,

ollama_classifier-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: uv 0.10.9
+Root-Is-Purelib: true
+Tag: py3-none-any