adapt classifier to outlines v1.0.4

classifier/classifier.py

@@ -2,11 +2,12 @@ import logging
 import os
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
-from typing import List, Optional
+from typing import List
+import json
 
 import outlines
 from outlines.models import openai
-from outlines.generate import choice
+from outlines.types import Choice
 
 # Configure logger
 logging.basicConfig(level=logging.DEBUG)
@@ -15,7 +16,9 @@ logger = logging.getLogger("classifier")
 app = FastAPI()
 
 # Global variables for shared config and classifier
-clf = None
+llm = None
+clf_output = None
+prompt_template = None
 config_set = False
 
 class Config(BaseModel):
@@ -33,7 +36,7 @@ class Resp(BaseModel):
 @app.post("/config")
 def configure(req: Config):
     """Receive and initialize classifier configuration."""
-    global clf, config_set
+    global llm, clf_output, prompt_template, config_set
 
     if config_set:
         logger.warning("Classifier already configured. Ignoring new config.")
@@ -43,11 +46,12 @@ def configure(req: Config):
     logger.debug(f"Using API_KEY: {'set' if api_key else 'missing'}")
 
     try:
+        # Instantiate model and choice output type
         llm = openai(req.model_name, api_key=api_key, base_url=req.base_url)
-        clf = choice(llm, req.class_set)
-        clf.class_set = req.class_set
-        clf.prompt_template = req.prompt_template
+        clf_output = Choice(req.class_set)
+        prompt_template = req.prompt_template
         config_set = True
+
         logger.info("Classifier configured successfully.")
         return {"status": "configured"}
     except Exception as e:
@@ -56,24 +60,32 @@ def configure(req: Config):
 
 @app.post("/classify", response_model=Resp)
 def classify(req: Req):
-    global clf
-    if clf is None or not config_set:
+    """Run text classification using the configured LLM."""
+    global llm, clf_output, prompt_template, config_set
+
+    if not config_set or llm is None or clf_output is None:
         raise HTTPException(status_code=503, detail="Classifier not configured yet")
 
-    # Render the prompt using the template
+    # Render prompt
     try:
-        prompt = clf.prompt_template.replace("{message}", req.message)
+        prompt = prompt_template.replace("{message}", req.message)
         logger.debug(f"Rendered prompt: {prompt!r}")
     except Exception as e:
         logger.warning(f"Prompt rendering failed: {e}")
         prompt = req.message
 
-    # Run classifier
+    # Invoke LLM classifier
     try:
-        result =  clf(prompt)
+        raw = llm(prompt, clf_output)
+        try:
+            # llm returned '{"result": "Label"}' → parse it
+            unwrapped = json.loads(raw).get("result", raw)
+        except Exception:
+            unwrapped = raw
+        result = unwrapped
         logger.debug(f"Classification result: {result}")
     except Exception as e:
-        logger.error(f"Classification error: {e}. Falling back to: {clf.class_set[-1]}")
-        result = clf.class_set[-1]
+        logger.error(f"Classification error: {e}. Falling back to: {req.class_set[-1]}")
+        result = req.class_set[-1]
 
     return Resp(result=result)

classifier/requirements.txt

@@ -1,6 +1,7 @@
 fastapi
 uvicorn
-outlines==0.2.1
+outlines
+numpy
 
 # Force CPU-only torch
 torch==2.0.1+cpu