changes to app.py

app.py

@@ -16,7 +16,7 @@ import fitz
 from PIL import Image
 import pytesseract
 from sentence_transformers import SentenceTransformer
-from ctransformers import AutoModel
+from ctransformers import AutoModelForCausalLM  # ✅ FIXED import
 
 # --- THIS IS THE FIX FOR TESSERACT ---
 # Explicitly tell pytesseract where to find the Tesseract OCR engine.
@@ -38,14 +38,18 @@ app.add_middleware(
 # --- Load Optimized Models ---
 try:
     logger.info("Loading optimized AI models...")
+
     # Using a smaller, but still powerful, BGE model
     embedding_model = SentenceTransformer('BAAI/bge-base-en-v1.5')
     
-    # Using TinyLlama, which is fast and efficient for CPU
-    llm = AutoModel.from_pretrained(
-        "TheBloke/TinyLlama-1.1B-Chat-v1.0-GGUF", 
-        model_file="tinyllama-1.1b-chat-v1.0.Q4_K_M.gguf"
+    # Load TinyLlama in GGUF format using ctransformers
+    llm = AutoModelForCausalLM.from_pretrained(
+        "TheBloke/TinyLlama-1.1B-Chat-v1.0-GGUF",
+        model_file="tinyllama-1.1b-chat-v1.0.Q4_K_M.gguf",
+        model_type="llama",   # Tell ctransformers the model family
+        gpu_layers=0          # For CPU-only environment
     )
+
     logger.info("AI models loaded successfully.")
 except Exception as e:
     logger.critical(f"Fatal error: Could not load AI models. {e}")
@@ -55,43 +59,62 @@ except Exception as e:
 SESSION_DATA = {}
 
 # --- 2. DATA MODELS ---
-class QueryRequest(BaseModel): question: str
-class UploadResponse(BaseModel): session_id: str; filename: str; chunks_created: int
-class QueryResponse(BaseModel): answer: str; context: str
+class QueryRequest(BaseModel):
+    question: str
+
+class UploadResponse(BaseModel):
+    session_id: str
+    filename: str
+    chunks_created: int
+
+class QueryResponse(BaseModel):
+    answer: str
+    context: str
 
 # --- 3. HELPER FUNCTIONS ---
 def parse_pdf(content: bytes) -> str:
-    doc = fitz.open(stream=content, filetype="pdf"); return "".join(page.get_text() for page in doc)
+    doc = fitz.open(stream=content, filetype="pdf")
+    return "".join(page.get_text() for page in doc)
 
 def parse_image(content: bytes) -> str:
-    image = Image.open(io.BytesIO(content)); return pytesseract.image_to_string(image)
+    image = Image.open(io.BytesIO(content))
+    return pytesseract.image_to_string(image)
 
 # --- 4. API ENDPOINTS ---
 
 @app.get("/")
-def read_root(): return {"status": "ok", "message": "Welcome to the Optimized Universal Data AI"}
+def read_root():
+    return {"status": "ok", "message": "Welcome to the Optimized Universal Data AI"}
 
 @app.post("/upload", response_model=UploadResponse)
 async def upload_file(file: UploadFile = File(...)):
-    if not embedding_model: raise HTTPException(status_code=503, detail="Embedding model not available.")
+    if not embedding_model:
+        raise HTTPException(status_code=503, detail="Embedding model not available.")
     
     session_id = str(uuid.uuid4())
     content = await file.read()
     content_type = file.content_type
     
-    if content_type == "application/pdf": text = parse_pdf(content)
-    elif content_type and content_type.startswith("image/"): text = parse_image(content)
-    elif file.filename.endswith(('.txt', '.md')): text = content.decode("utf-8")
-    else: raise HTTPException(status_code=400, detail=f"Unsupported file type: {content_type}")
-
-    if not text.strip(): raise HTTPException(status_code=400, detail="No text could be extracted.")
+    if content_type == "application/pdf":
+        text = parse_pdf(content)
+    elif content_type and content_type.startswith("image/"):
+        text = parse_image(content)
+    elif file.filename.endswith(('.txt', '.md')):
+        text = content.decode("utf-8")
+    else:
+        raise HTTPException(status_code=400, detail=f"Unsupported file type: {content_type}")
+
+    if not text.strip():
+        raise HTTPException(status_code=400, detail="No text could be extracted.")
     
     text_chunks = semantic_chunker(text, embedding_model)
-    if not text_chunks: raise HTTPException(status_code=400, detail="Document too short to be processed.")
+    if not text_chunks:
+        raise HTTPException(status_code=400, detail="Document too short to be processed.")
     
     embeddings = embedding_model.encode(text_chunks, convert_to_numpy=True)
     serialized_index = create_faiss_index(embeddings)
-    if not serialized_index: raise HTTPException(status_code=500, detail="Failed to create document index.")
+    if not serialized_index:
+        raise HTTPException(status_code=500, detail="Failed to create document index.")
     
     SESSION_DATA[session_id] = {"chunks": text_chunks, "index": serialized_index}
     logger.info(f"Session {session_id} created with {len(text_chunks)} chunks.")
@@ -110,7 +133,8 @@ async def query_session(session_id: str, request: QueryRequest):
     question_embedding = embedding_model.encode([query_with_prefix], convert_to_numpy=True).astype('float32')
     
     index = deserialize_faiss_index(session["index"])
-    if not index: raise HTTPException(status_code=500, detail="Could not load session index.")
+    if not index:
+        raise HTTPException(status_code=500, detail="Could not load session index.")
     
     k = min(5, index.ntotal)
     distances, indices = index.search(question_embedding, k)
@@ -135,4 +159,4 @@ Question: {request.question}<|im_end|>
         stop=["<|im_end|>"]
     )
     
-    return {"answer": answer.strip(), "context": context}
+    return {"answer": answer.strip(), "context": context}