push

app.py

@@ -10,7 +10,7 @@ from typing import Optional, Tuple
 from fastapi import FastAPI, UploadFile, File, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import JSONResponse
-from transformers import pipeline
+from transformers import AutoTokenizer, AutoModelForCausalLM
 from docx import Document as DocxDocument
 from pptx import Presentation
 import logging
@@ -44,14 +44,15 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-MODEL_ID = "tiiuae/Falcon3-3B-Instruct"
-pipe = None
+MODEL_ID = "HuggingFaceH4/zephyr-7b-beta"
+model = None
+tokenizer = None
 ocr_reader = None
 
 @app.on_event("startup")
 async def load_model():
-    """Load the model pipeline and OCR reader on startup"""
-    global pipe, ocr_reader
+    """Load the Zephyr tokenizer/model and OCR reader on startup"""
+    global tokenizer, model, ocr_reader
     try:
         logger.info(f"Loading model: {MODEL_ID} ...")
         logger.info("Optimizing for CPU-only inference...")
@@ -60,19 +61,20 @@ async def load_model():
         torch.set_num_interop_threads(os.cpu_count() or 4)
         
         logger.info(f"Using {torch.get_num_threads()} CPU threads for inference")
-        logger.info("Loading full model into CPU RAM (no offloading)...")
-        
-        pipe = pipeline(
-            "text-generation",
-            model=MODEL_ID,
-            dtype=torch.bfloat16,
+        logger.info("Loading Zephyr tokenizer and model (CPU)...")
+
+        tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+        if tokenizer.pad_token is None:
+            tokenizer.pad_token = tokenizer.eos_token
+
+        desired_dtype = torch.bfloat16 if hasattr(torch, "cpu") and getattr(torch.cpu, "is_bf16_supported", lambda: False)() else torch.float32
+        model = AutoModelForCausalLM.from_pretrained(
+            MODEL_ID,
+            dtype=desired_dtype,
             device_map="cpu",
-            model_kwargs={
-                "low_cpu_mem_usage": False,
-                "offload_folder": None
-            }
-        )
-        logger.info("✅ Model loaded successfully in CPU RAM!")
+            low_cpu_mem_usage=False
+        ).eval()
+        logger.info("✅ Zephyr loaded successfully on CPU!")
         
         logger.info("Loading OCR reader...")
         try:
@@ -437,68 +439,46 @@ Produce ONLY valid JSON with these exact fields:
 }}"""
 
     try:
-        full_prompt = f"{system_message}\n\n{user_message}"
-        
-        logger.info(f"Input prompt length: {len(full_prompt)} characters")
-        logger.info("Starting model generation with pipeline...")
-        
-        start_time = time.time()
-        
         messages = [
-            {"role": "user", "content": full_prompt}
+            {"role": "system", "content": system_message},
+            {"role": "user", "content": user_message}
         ]
-        
-        result = pipe(
-            messages,
-            max_new_tokens=800,
-            temperature=0.2,
-            do_sample=True,
-            top_p=0.9,
-            return_full_text=False
-        )
-        
+
+        start_time = time.time()
+        prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+        inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=3800).to(model.device)
+
+        with torch.no_grad():
+            outputs = model.generate(
+                **inputs,
+                max_new_tokens=800,
+                temperature=0.2,
+                top_p=0.9,
+                do_sample=True,
+                repetition_penalty=1.08,
+                pad_token_id=tokenizer.eos_token_id,
+                use_cache=True
+            )
+
         generation_time = time.time() - start_time
-        raw_output = result[0]["generated_text"]
-        logger.info(f"✅ Generated {len(raw_output)} characters in {generation_time:.2f}s ({len(raw_output)/generation_time:.1f} chars/sec)")
-        
-        start = raw_output.find('{')
-        end = raw_output.rfind('}') + 1
-        
-        if start == -1 or end == 0:
-            logger.warning("No JSON found in output, returning raw output")
-            raise ValueError(f"No JSON object found in model output. Raw output: {raw_output[:500]}")
-        
-        json_str = raw_output[start:end]
-        
+        raw_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        if "<|assistant|>" in raw_text:
+            raw_text = raw_text.split("<|assistant|>")[-1]
+
+        logger.info(f"✅ Generated {len(raw_text)} chars in {generation_time:.2f}s")
+
+        start = raw_text.find('{')
+        end = raw_text.rfind('}') + 1
+        if start == -1 or end <= 0:
+            raise ValueError("No JSON object found in model output")
+
+        json_str = raw_text[start:end]
         try:
-            parsed_json = json.loads(json_str)
-            return parsed_json
+            return json.loads(json_str)
         except json.JSONDecodeError as e:
-            logger.warning(f"JSON parsing failed, attempting to repair: {e}")
-            try:
-                repaired_json = _repair_json(json_str)
-                parsed_json = json.loads(repaired_json)
-                logger.info("✅ JSON successfully repaired")
-                return parsed_json
-            except Exception as repair_error:
-                logger.error(f"JSON repair also failed: {repair_error}")
-                logger.error(f"Problematic JSON (around error): {json_str[max(0, e.pos-200):e.pos+200]}")
-                
-                try:
-                    import json5
-                    parsed_json = json5.loads(json_str)
-                    logger.info("✅ JSON5 parsing succeeded as fallback")
-                    return parsed_json
-                except ImportError:
-                    pass
-                except Exception:
-                    pass
-                
-                raise ValueError(f"Failed to parse JSON from model output at position {e.pos}: {str(e)}. JSON preview: {json_str[max(0, e.pos-200):e.pos+200]}")
-        
-    except json.JSONDecodeError as e:
-        logger.error(f"JSON parsing error: {e}")
-        raise ValueError(f"Failed to parse JSON from model output: {str(e)}")
+            repaired = _repair_json(json_str)
+            return json.loads(repaired)
+
     except Exception as e:
         logger.error(f"Model generation error: {e}")
         raise ValueError(f"Error during model inference: {str(e)}")
@@ -522,19 +502,28 @@ Full Deck Length: {len(full_text)} characters
 Produce a FINAL comprehensive review with the same JSON structure as before, consolidating all findings."""
     
     try:
-        full_prompt = f"{system_message}\n\n{user_message}"
-        messages = [{"role": "user", "content": full_prompt}]
-        
-        result = pipe(
-            messages,
-            max_new_tokens=800,
-            temperature=0.2,
-            do_sample=True,
-            top_p=0.9,
-            return_full_text=False
-        )
-        
-        raw_output = result[0]["generated_text"]
+        messages = [
+            {"role": "system", "content": system_message},
+            {"role": "user", "content": user_message}
+        ]
+        prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+        inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=3800).to(model.device)
+
+        with torch.no_grad():
+            outputs = model.generate(
+                **inputs,
+                max_new_tokens=800,
+                temperature=0.2,
+                top_p=0.9,
+                do_sample=True,
+                repetition_penalty=1.05,
+                pad_token_id=tokenizer.eos_token_id,
+                use_cache=True
+            )
+
+        raw_output = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        if "<|assistant|>" in raw_output:
+            raw_output = raw_output.split("<|assistant|>")[-1]
         
         start = raw_output.find('{')
         end = raw_output.rfind('}') + 1
@@ -554,11 +543,9 @@ Produce a FINAL comprehensive review with the same JSON structure as before, con
         
         try:
             combined_json = json.loads(json_str)
-        except json.JSONDecodeError as e:
-            logger.warning(f"JSON parsing failed in combine, attempting repair: {e}")
+        except json.JSONDecodeError:
             try:
-                repaired_json = _repair_json(json_str)
-                combined_json = json.loads(repaired_json)
+                combined_json = json.loads(_repair_json(json_str))
             except Exception:
                 logger.warning("JSON repair failed, returning basic structure")
                 return {
@@ -641,21 +628,28 @@ Return ONLY valid JSON:
 }}"""
 
     try:
-        full_prompt = f"{system_message}\n\n{user_message}"
-        messages = [{"role": "user", "content": full_prompt}]
-        
-        result = pipe(
-            messages,
-            max_new_tokens=600,
-            temperature=0.25,
-            do_sample=True,
-            top_p=0.9,
-            return_full_text=False,
-            pad_token_id=None,
-            eos_token_id=None
-        )
-        
-        raw_output = result[0]["generated_text"]
+        messages = [
+            {"role": "system", "content": system_message},
+            {"role": "user", "content": user_message}
+        ]
+        prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+        inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=3600).to(model.device)
+
+        with torch.no_grad():
+            outputs = model.generate(
+                **inputs,
+                max_new_tokens=600,
+                temperature=0.25,
+                top_p=0.9,
+                do_sample=True,
+                repetition_penalty=1.05,
+                pad_token_id=tokenizer.eos_token_id,
+                use_cache=True
+            )
+
+        raw_output = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        if "<|assistant|>" in raw_output:
+            raw_output = raw_output.split("<|assistant|>")[-1]
         
         start = raw_output.find('{')
         end = raw_output.rfind('}') + 1
@@ -670,11 +664,9 @@ Return ONLY valid JSON:
         
         try:
             improvement_json = json.loads(json_str)
-        except json.JSONDecodeError as e:
-            logger.warning(f"JSON parsing failed in improvements, attempting repair: {e}")
+        except json.JSONDecodeError:
             try:
-                repaired_json = _repair_json(json_str)
-                improvement_json = json.loads(repaired_json)
+                improvement_json = json.loads(_repair_json(json_str))
             except Exception:
                 logger.warning("JSON repair failed, returning default improvement structure")
                 return {
@@ -707,7 +699,7 @@ async def health():
     """Health check endpoint"""
     return {
         "status": "healthy",
-        "model_loaded": pipe is not None
+        "model_loaded": (model is not None and tokenizer is not None)
     }
 
 @app.post("/review")
@@ -717,17 +709,26 @@ async def review_deck(file: UploadFile = File(...)):
     
     Supported formats: PDF, DOCX, PPT, PPTX
     """
-    if pipe is None:
-        raise HTTPException(status_code=503, detail="Model not loaded yet. Please wait for startup to complete.")
-    
-    file_extension = Path(file.filename).suffix.lower()
-    supported_extensions = [".pdf", ".docx", ".doc", ".ppt", ".pptx"]
-    
-    if file_extension not in supported_extensions:
-        raise HTTPException(
-            status_code=400,
-            detail=f"Unsupported file type: {file_extension}. Supported: {', '.join(supported_extensions)}"
-        )
+    try:
+        if model is None or tokenizer is None:
+            raise HTTPException(status_code=503, detail="Model not loaded yet. Please wait for startup to complete.")
+        
+        if not file.filename:
+            raise HTTPException(status_code=400, detail="Filename is missing")
+        
+        file_extension = Path(file.filename).suffix.lower()
+        supported_extensions = [".pdf", ".docx", ".doc", ".ppt", ".pptx"]
+        
+        if file_extension not in supported_extensions:
+            raise HTTPException(
+                status_code=400,
+                detail=f"Unsupported file type: {file_extension}. Supported: {', '.join(supported_extensions)}"
+            )
+    except HTTPException:
+        raise
+    except Exception as e:
+        logger.error(f"Error in request validation: {e}", exc_info=True)
+        raise HTTPException(status_code=500, detail=f"Request validation error: {str(e)}")
     
     temp_file = None
     try:
@@ -758,20 +759,29 @@ async def review_deck(file: UploadFile = File(...)):
             logger.info("Review generated successfully")
             
             logger.info("Checking if improvement pointers are needed...")
-            improvement_pointers = generate_improvement_pointers(review_result)
-            review_result["improvement_analysis"] = improvement_pointers
+            try:
+                improvement_pointers = generate_improvement_pointers(review_result)
+                review_result["improvement_analysis"] = improvement_pointers
+            except Exception as imp_error:
+                logger.warning(f"Improvement pointers generation failed: {imp_error}, continuing without it")
+                review_result["improvement_analysis"] = {
+                    "needs_improvement": True,
+                    "improvement_pointers": [],
+                    "error": "Failed to generate improvement pointers"
+                }
             
             return JSONResponse(content=review_result)
         except ValueError as e:
+            logger.error(f"ValueError in review generation: {e}", exc_info=True)
             raise HTTPException(status_code=500, detail=str(e))
         except Exception as e:
-            logger.error(f"Review generation error: {e}")
+            logger.error(f"Review generation error: {e}", exc_info=True)
             raise HTTPException(status_code=500, detail=f"Error generating review: {str(e)}")
     
     except HTTPException:
         raise
     except Exception as e:
-        logger.error(f"Unexpected error: {e}")
+        logger.error(f"Unexpected error in review endpoint: {e}", exc_info=True)
         raise HTTPException(status_code=500, detail=f"Unexpected error: {str(e)}")
     
     finally: