ChintanSatva commited on
Commit
18ba589
·
verified ·
1 Parent(s): 256943c

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +14 -10
app.py CHANGED
@@ -1,4 +1,4 @@
1
- from fastapi import FastAPI, File, UploadFile, HTTPException
2
  import pytesseract
3
  import cv2
4
  import os
@@ -18,25 +18,29 @@ import cachetools
18
  import hashlib
19
  from vllm import LLM
20
 
21
- app = FastAPI()
22
 
23
  # Configure logging
24
- logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
25
  logger = logging.getLogger(__name__)
26
 
27
  # Set Tesseract path
28
- pytesseract.pytesseract.tesseract_cmd = "/usr/bin/tesseract"
29
 
30
- # Initialize BitNet model
31
  try:
32
- llm = LLM(model="bitnet/BitNet-b1.2-3B", gpu_memory_utilization=0.0) # CPU-only
 
 
 
 
33
  except Exception as e:
34
  logger.error(f"Failed to load BitNet model: {str(e)}")
35
- raise HTTPException(status_code=500, detail="BitNet model initialization failed")
36
 
37
- # In-memory caches (1-hour TTL)
38
- raw_text_cache = cachetools.TTLCache(maxsize=100, ttl=3600)
39
- structured_data_cache = cachetools.TTLCache(maxsize=100, ttl=3600)
40
 
41
  def log_memory_usage():
42
  """Log current memory usage."""
 
1
+ from fastapi import FastAPI, File, UploadFile
2
  import pytesseract
3
  import cv2
4
  import os
 
18
  import hashlib
19
  from vllm import LLM
20
 
21
+ app = LLMAPI()
22
 
23
  # Configure logging
24
+ logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(name)s - %(message)s')
25
  logger = logging.getLogger(__name__)
26
 
27
  # Set Tesseract path
28
+ pytesseract.pytesseract_cmd = "/usr/bin/tesseract"
29
 
30
+ # Initialize BitNet model for CPU-only
31
  try:
32
+ llm = LLM(
33
+ model="username/bitnet-finetuned-invoice", # Replace with your fine-tuned BitNet model
34
+ device="cuda", # Explicitly set to CPU
35
+ enforce_eager=True, # Disable CUDA graph compilation for CPU
36
+ )
37
  except Exception as e:
38
  logger.error(f"Failed to load BitNet model: {str(e)}")
39
+ raise RuntimeError("Failed to initialize BitNet model")
40
 
41
+ # In-memory caches (1GB TTL)
42
+ raw_text_cache = cachetools.TTLTTL(maxsize=100, ttl=3600)
43
+ structured_data_cache = cachetools.TTLTTL(maxsize=100, ttl=3600)
44
 
45
  def log_memory_usage():
46
  """Log current memory usage."""