Update app.py

app.py

@@ -1,6 +1,5 @@
 import os
 import re
-import asyncio
 import time
 from typing import Literal, Optional
 from fastapi import FastAPI, HTTPException
@@ -10,14 +9,12 @@ from llama_cpp import Llama
 from contextlib import asynccontextmanager
 
 # ==================== CONFIGURATION ====================
-# Hugging Face Spaces optimized settings
 MODEL_REPO = "bartowski/Phi-3.1-mini-4k-instruct-GGUF"
 MODEL_FILE = "Phi-3.1-mini-4k-instruct-IQ2_M.gguf"
-MODEL_PATH = os.environ.get("MODEL_PATH", MODEL_FILE)
 
-# CPU settings optimized for Spaces (2 CPU cores typical)
+# Hugging Face Spaces optimized settings
 N_THREADS = int(os.environ.get("N_THREADS", "2"))
-N_CTX = int(os.environ.get("N_CTX", "2048"))  # Reduced for faster inference
+N_CTX = int(os.environ.get("N_CTX", "2048"))  # Reduced for faster loading
 N_BATCH = int(os.environ.get("N_BATCH", "128"))
 N_GPU_LAYERS = int(os.environ.get("N_GPU_LAYERS", "0"))
 
@@ -26,20 +23,34 @@ END_TOKEN = "<|endoftext|>"
 
 # ==================== GLOBAL MODEL ====================
 llm = None
+model_loading_error = None
 
 # ==================== LIFECYCLE MANAGEMENT ====================
 @asynccontextmanager
 async def lifespan(app: FastAPI):
-    # Startup
-    print("🚀 Starting FormatAI Humanizer Backend on Hugging Face Space")
-    print(f"📊 Configuration:")
-    print(f"   Model: {MODEL_PATH}")
-    print(f"   Threads: {N_THREADS}")
-    print(f"   Context: {N_CTX}")
+    # Startup - load model
+    global llm, model_loading_error
+    print("🚀 Starting FormatAI Humanizer Backend")
+    print(f"📊 Configuration: Threads={N_THREADS}, Context={N_CTX}")
     
-    # Load model on startup
-    global llm
-    llm = load_model()
+    try:
+        print(f"📥 Downloading model: {MODEL_REPO}/{MODEL_FILE}")
+        llm = Llama.from_pretrained(
+            repo_id=MODEL_REPO,
+            filename=MODEL_FILE,
+            n_threads=N_THREADS,
+            n_ctx=N_CTX,
+            n_batch=N_BATCH,
+            n_gpu_layers=N_GPU_LAYERS,
+            verbose=True,  # Set to True to see loading progress
+            use_mlock=False,  # Important for Spaces
+        )
+        print("✅ Model loaded successfully!")
+        model_loading_error = None
+    except Exception as e:
+        print(f"❌ Model loading failed: {e}")
+        model_loading_error = str(e)
+        llm = None
     
     yield
     
@@ -52,16 +63,14 @@ async def lifespan(app: FastAPI):
 app = FastAPI(
     title="FormatAI Humanizer API",
     description="Backend API for text transformation with Phi-3.1 Mini",
-    version="1.0.0",
+    version="2.0.0",
     lifespan=lifespan
 )
 
 # CORS - Allow your Vercel frontend
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=[
-        "*",  # Allow all origins temporarily, update with your Vercel URL
-    ],
+    allow_origins=["*"],
     allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
@@ -77,7 +86,8 @@ class HumanizeRequest(BaseModel):
 
 # ==================== STYLE PROMPTS ====================
 STYLE_PROMPTS = {
-    "professional": """You are a professional writing assistant. Rewrite the text below in formal, corporate business language.
+    "professional": """<|system|>
+You are a professional writing assistant. Rewrite the text below in formal, corporate business language.
 
 IMPORTANT RULES:
 1. Output ONLY the rewritten text
@@ -85,12 +95,15 @@ IMPORTANT RULES:
 3. Keep the same meaning
 4. Use formal vocabulary
 5. Proper grammar and structure
-
+<|end|>
+<|user|>
 Text to rewrite: {text}
 
-Rewritten (professional):""",
+Rewritten (professional):<|end|>
+<|assistant|>""",
 
-    "casual": """You are a casual writing assistant. Rewrite the text below in friendly, natural, conversational English.
+    "casual": """<|system|>
+You are a casual writing assistant. Rewrite the text below in friendly, natural, conversational English.
 
 IMPORTANT RULES:
 1. Output ONLY the rewritten text
@@ -98,12 +111,15 @@ IMPORTANT RULES:
 3. Keep the same meaning
 4. Use contractions (I'm, don't, etc.)
 5. Sound like a real person speaking
-
+<|end|>
+<|user|>
 Text to rewrite: {text}
 
-Rewritten (casual):""",
+Rewritten (casual):<|end|>
+<|assistant|>""",
 
-    "academic": """You are an academic writing assistant. Rewrite the text below in formal scholarly language.
+    "academic": """<|system|>
+You are an academic writing assistant. Rewrite the text below in formal scholarly language.
 
 IMPORTANT RULES:
 1. Output ONLY the rewritten text
@@ -111,12 +127,15 @@ IMPORTANT RULES:
 3. Keep the same meaning
 4. Use precise academic vocabulary
 5. Maintain formal structure
-
+<|end|>
+<|user|>
 Text to rewrite: {text}
 
-Rewritten (academic):""",
+Rewritten (academic):<|end|>
+<|assistant|>""",
 
-    "marketing": """You are a marketing copywriter. Rewrite the text below into persuasive marketing language.
+    "marketing": """<|system|>
+You are a marketing copywriter. Rewrite the text below into persuasive marketing language.
 
 IMPORTANT RULES:
 1. Output ONLY the rewritten text
@@ -124,10 +143,12 @@ IMPORTANT RULES:
 3. Keep the same meaning
 4. Use emotional hooks and benefits
 5. Make it engaging and compelling
-
+<|end|>
+<|user|>
 Text to rewrite: {text}
 
-Rewritten (marketing):"""
+Rewritten (marketing):<|end|>
+<|assistant|>"""
 }
 
 STYLE_TEMPERATURES = {
@@ -138,68 +159,17 @@ STYLE_TEMPERATURES = {
 }
 
 # ==================== HELPER FUNCTIONS ====================
-def load_model():
-    """Load the GGUF model"""
-    global MODEL_PATH  # Access the global MODEL_PATH
-    
-    print(f"🔄 Loading model from: {MODEL_PATH}")
-    
-    try:
-        # Check if model exists locally
-        if not os.path.exists(MODEL_PATH):
-            print("📥 Downloading model from Hugging Face Hub...")
-            try:
-                from huggingface_hub import hf_hub_download
-                # Use a different variable name to avoid conflict
-                downloaded_path = hf_hub_download(
-                    repo_id=MODEL_REPO,
-                    filename=MODEL_FILE,
-                    local_dir=".",
-                    token=os.environ.get("HF_TOKEN", None)
-                )
-                print(f"✅ Model downloaded to: {downloaded_path}")
-            except ImportError:
-                print("⚠️  huggingface-hub not installed, using local model path")
-                # If we can't download, use fallback path
-                MODEL_PATH = os.path.join("/code", MODEL_FILE)
-        
-        print(f"📁 Model path: {MODEL_PATH}")
-        
-        model = Llama(
-            model_path=MODEL_PATH,
-            n_threads=N_THREADS,
-            n_ctx=N_CTX,
-            n_batch=N_BATCH,
-            n_gpu_layers=N_GPU_LAYERS,
-            verbose=False,
-            use_mlock=False,  # Important for Spaces
-        )
-        
-        print(f"✅ Model loaded successfully!")
-        return model
-        
-    except Exception as e:
-        print(f"❌ Failed to load model: {e}")
-        import traceback
-        traceback.print_exc()
-        return None
-
 def clean_output(text: str) -> str:
     """Clean model output"""
     if not text:
         return ""
     
-    # Remove common artifacts
-    clean = re.sub(r'Rewritten\s*\([^)]+\):', '', text, flags=re.IGNORECASE)
-    clean = re.sub(r'IMPORTANT RULES:.*?(?=\n\n|\Z)', '', clean, flags=re.DOTALL)
-    clean = re.sub(r'You are [^\.]+\.', '', clean)
-    
     # Remove Phi-3.1 special tokens
-    clean = re.sub(r'<\|[^>]+\|>', '', clean)
-    clean = re.sub(r'\[/?[^]]+\]', '', clean)
+    clean = re.sub(r'<\|[^>]+\|>', '', text)
+    clean = re.sub(r'Rewritten\s*\([^)]+\):', '', clean, flags=re.IGNORECASE)
+    clean = re.sub(r'Text to rewrite:.*', '', clean, flags=re.DOTALL)
     
     # Clean whitespace
-    clean = re.sub(r'\n+', ' ', clean)
     clean = re.sub(r'\s+', ' ', clean)
     clean = clean.strip()
     
@@ -211,24 +181,16 @@ def clean_output(text: str) -> str:
 
 def format_prompt(text: str, style: str) -> str:
     """Format prompt for Phi-3.1"""
-    system_prompt = STYLE_PROMPTS[style].format(text=text)
-    
-    # Phi-3.1 chat format
-    prompt = f"<|system|>\n{system_prompt}\n<|end|>\n"
-    prompt += f"<|user|>\nPlease rewrite this text in {style} style:\n{text}\n<|end|>\n"
-    prompt += "<|assistant|>\n"
-    
-    return prompt
+    return STYLE_PROMPTS[style].format(text=text)
 
 async def transform_with_model(text: str, style: str) -> str:
     """Transform text using the loaded model"""
-    global llm
+    global llm, model_loading_error
     
     if llm is None:
-        # Try to load model if not loaded
-        llm = load_model()
-        if llm is None:
-            raise HTTPException(status_code=503, detail="Model not available. Please check if model file exists.")
+        if model_loading_error:
+            raise HTTPException(status_code=503, detail=f"Model failed to load: {model_loading_error}")
+        raise HTTPException(status_code=503, detail="Model not available. Please wait for model to load.")
     
     try:
         # Build prompt
@@ -240,7 +202,7 @@ async def transform_with_model(text: str, style: str) -> str:
         
         output = llm(
             prompt,
-            max_tokens=min(400, len(text) + 100),  # Dynamic token limit
+            max_tokens=min(400, len(text) + 100),
             temperature=temperature,
             top_p=0.9,
             repeat_penalty=1.1,
@@ -263,13 +225,11 @@ async def transform_with_model(text: str, style: str) -> str:
         if not cleaned or cleaned.isspace():
             cleaned = f"[{style.capitalize()} Version]: {text}"
         
-        print(f"✅ Transformation completed in {processing_time:.2f}s")
+        print(f"✅ Transformation completed in {processing_time:.2f}s (Style: {style})")
         return cleaned
         
     except Exception as e:
         print(f"❌ Model error: {e}")
-        import traceback
-        traceback.print_exc()
         raise HTTPException(status_code=500, detail=f"Model error: {str(e)}")
 
 # ==================== API ENDPOINTS ====================
@@ -277,9 +237,11 @@ async def transform_with_model(text: str, style: str) -> str:
 async def root():
     """Health check endpoint"""
     return {
-        "status": "online",
+        "status": "online" if llm else "model_loading_failed",
         "service": "FormatAI Humanizer",
-        "model": "Phi-3.1-mini-4k-instruct-GGUF",
+        "model": MODEL_FILE,
+        "model_loaded": llm is not None,
+        "model_error": model_loading_error,
         "styles_available": list(STYLE_PROMPTS.keys()),
         "max_input_length": MAX_INPUT_LENGTH
     }
@@ -288,11 +250,11 @@ async def root():
 async def health_check():
     """Detailed health check"""
     return {
-        "status": "healthy" if llm else "model_loading_failed",
+        "status": "healthy" if llm else "unhealthy",
         "model_loaded": llm is not None,
+        "model_error": model_loading_error,
         "threads": N_THREADS,
-        "context_size": N_CTX,
-        "model_path": MODEL_PATH
+        "context_size": N_CTX
     }
 
 @app.post("/api/humanize")
@@ -357,8 +319,6 @@ async def transform_text(request: TransformRequest):
         raise
     except Exception as e:
         print(f"❌ Transformation error: {e}")
-        import traceback
-        traceback.print_exc()
         raise HTTPException(status_code=500, detail=f"Transformation failed: {str(e)}")
 
 @app.get("/api/styles")
@@ -367,13 +327,12 @@ async def get_styles():
     styles_info = {}
     for style, prompt in STYLE_PROMPTS.items():
         # Extract first line for description
-        first_line = prompt.split('\n')[0]
-        description = first_line.replace("You are ", "").replace(".", "")
+        first_line = prompt.split('\n')[0].replace('<|system|>', '').strip()
         
         styles_info[style] = {
-            "description": description,
+            "description": first_line,
             "temperature": STYLE_TEMPERATURES[style],
-            "example_prompt": prompt[:100] + "..." if len(prompt) > 100 else prompt
+            "max_tokens": 400
         }
     
     return {
@@ -391,5 +350,5 @@ if __name__ == "__main__":
         "app:app",
         host="0.0.0.0",
         port=port,
-        reload=False  # Disable reload for production
+        reload=False
     )