Update app.py

app.py

@@ -3,13 +3,14 @@ import logging
 import json
 import os
 from pydantic import BaseModel
-from vllm import LLM
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+import torch
 import psutil
 import cachetools
 import hashlib
 
-# Set environment variable for transformers cache
-os.environ["TRANSFORMERS_CACHE"] = "/app/cache"
+# Set environment variable for cache
+os.environ["HF_HOME"] = "/app/cache"
 
 app = FastAPI()
 
@@ -17,21 +18,17 @@ app = FastAPI()
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 
-# Set cache directories
-HF_HOME = "/app/cache"
-NUMBA_CACHE_DIR = "/app/cache"
-
-# Initialize BitNet model for CPU-only
+# Initialize BitNet model and tokenizer
 try:
-    llm = LLM(
-        model="1bitLLM/bitnet_b1_58-3B",
-        device="cpu",
-        enforce_eager=True,
-        tensor_parallel_size=1,
-        disable_custom_all_reduce=True,
-        max_model_len=2048,
-        dtype="float32",
+    model_name = "1bitLLM/bitnet_b1_58-3B"
+    tokenizer = AutoTokenizer.from_pretrained(model_name, cache_dir="/app/cache")
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name,
+        torch_dtype=torch.float32,
+        device_map="cpu",
+        cache_dir="/app/cache"
     )
+    pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, max_new_tokens=512)
 except Exception as e:
     logger.error(f"Failed to load BitNet model: {str(e)}")
     raise HTTPException(status_code=500, detail=f"BitNet model initialization failed: {str(e)}")
@@ -107,11 +104,10 @@ Output JSON:
   "subcategory_confidence": 0.0
 }}
 """
-        outputs = llm.generate(prompts=[prompt])
-        json_str = outputs[0].outputs[0].text
-        json_start = json_str.find("{")
-        json_end = json_str.rfind("}") + 1
-        result = json.loads(json_str[json_start:json_end])
+        outputs = pipe(prompt)[0]["generated_text"]
+        json_start = outputs.rfind("{")
+        json_end = outputs.rfind("}") + 1
+        result = json.loads(outputs[json_start:json_end])
 
         # Normalize category and subcategory
         def normalize(s):