Update app.py

app.py

@@ -1,185 +1,393 @@
-import asyncio
+import os
 import re
-from typing import Literal
+import asyncio
+import time
+from typing import Literal, Optional
 from fastapi import FastAPI, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 from llama_cpp import Llama
-
-# ---------------- MODEL CONFIG ---------------- #
-
-# Model is in ROOT of the repository, not inside a folder
-MODEL_PATH = "/code/Phi-3.1-mini-4k-instruct-IQ2_M.gguf"
-
-# CPU settings for llama.cpp
-N_THREADS = 4
-N_CTX = 4096
-N_BATCH = 512
-N_GPU_LAYERS = 0
-
-# Concurrency limit
-MAX_CONCURRENT_REQUESTS = 6
-
-# Token to force stopping the generation
-END_TOKEN = "###END_OF_RESPONSE###"
-
-print("Loading model:", MODEL_PATH)
-llm = 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,
+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)
+N_THREADS = int(os.environ.get("N_THREADS", "2"))
+N_CTX = int(os.environ.get("N_CTX", "2048"))  # Reduced for faster inference
+N_BATCH = int(os.environ.get("N_BATCH", "128"))
+N_GPU_LAYERS = int(os.environ.get("N_GPU_LAYERS", "0"))
+
+MAX_INPUT_LENGTH = 1500
+END_TOKEN = "<|endoftext|>"
+
+# ==================== 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}")
+    
+    # Load model on startup
+    global llm
+    llm = load_model()
+    
+    yield
+    
+    # Shutdown
+    print("👋 Shutting down...")
+    if llm:
+        del llm
+
+# ==================== FASTAPI APP ====================
+app = FastAPI(
+    title="FormatAI Humanizer API",
+    description="Backend API for text transformation with Phi-3.1 Mini",
+    version="1.0.0",
+    lifespan=lifespan
 )
-print("Model loaded successfully.")
-
-# ---------------- FASTAPI APP ---------------- #
-
-app = FastAPI(title="FormatAI Humanizer Backend")
 
+# CORS - Allow your Vercel frontend
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],  # Vercel frontend allowed
+    allow_origins=[
+        "https://your-vercel-app.vercel.app",  # Replace with your actual Vercel URL
+        "http://localhost:3000",  # For local development
+        "http://localhost:5173",  # Vite dev server
+    ],
+    allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
 )
 
-# ---------------- REQUEST MODELS ---------------- #
-
+# ==================== DATA MODELS ====================
 class TransformRequest(BaseModel):
-    text: str
-    style: Literal["professional", "casual", "academic", "marketing"]
-
-class HumanizeRequest(BaseModel):  # legacy
-    text: str
+    text: str = Field(..., min_length=1, max_length=MAX_INPUT_LENGTH)
+    style: Literal["professional", "casual", "academic", "marketing"] = "casual"
 
+class HumanizeRequest(BaseModel):
+    text: str = Field(..., min_length=1, max_length=MAX_INPUT_LENGTH)
 
-# ---------------- STYLE PROMPTS ---------------- #
+# ==================== GLOBAL MODEL ====================
+llm = None
 
+# ==================== STYLE PROMPTS ====================
 STYLE_PROMPTS = {
-    "professional": (
-        "STYLE: PROFESSIONAL\n"
-        "Rewrite the user's text in a STRICTLY professional, corporate, formal tone. "
-        "Use respectful and clear business language. Do NOT add explanations. "
-        f"Output ONLY the rewritten text, then write {END_TOKEN}."
-    ),
-
-    "casual": (
-        "STYLE: CASUAL\n"
-        "Rewrite the user's text in a friendly, conversational, relaxed tone. "
-        "Use contractions and natural flow. Do NOT add explanations. "
-        f"Output ONLY the rewritten text, then write {END_TOKEN}."
-    ),
-
-    "academic": (
-        "STYLE: ACADEMIC\n"
-        "Rewrite the user's text in precise, formal academic language suitable for scholarly writing. "
-        "Use objective vocabulary and clear structure. Do NOT add explanations. "
-        f"Output ONLY the rewritten text, then write {END_TOKEN}."
-    ),
-
-    "marketing": (
-        "STYLE: MARKETING\n"
-        "Rewrite the user's text into persuasive, compelling marketing language. "
-        "Use emotional hooks, strong benefits, and engaging tone. Do NOT add explanations. "
-        f"Output ONLY the rewritten text, then write {END_TOKEN}."
-    ),
-}
-
+    "professional": """You are a professional writing assistant. Rewrite the text below in formal, corporate business language.
 
-# ---------------- HELPERS ---------------- #
+IMPORTANT RULES:
+1. Output ONLY the rewritten text
+2. No explanations, no notes
+3. Keep the same meaning
+4. Use formal vocabulary
+5. Proper grammar and structure
 
-def clean_output(raw: str) -> str:
-    """Trim unwanted tokens and cleanup."""
-    if not raw:
-        return ""
+Text to rewrite: {text}
 
-    # Remove system markers if any appear
-    raw = re.sub(r"<\|/?(system|assistant|user|end)\|>", "", raw, flags=re.I)
+Rewritten (professional):""",
 
-    # Cut off at END_TOKEN
-    if END_TOKEN in raw:
-        raw = raw.split(END_TOKEN)[0]
+    "casual": """You are a casual writing assistant. Rewrite the text below in friendly, natural, conversational English.
 
-    raw = raw.strip()
-    raw = re.sub(r"[ \t]+", " ", raw)
-    return raw.strip()
+IMPORTANT RULES:
+1. Output ONLY the rewritten text
+2. No explanations, no notes
+3. Keep the same meaning
+4. Use contractions (I'm, don't, etc.)
+5. Sound like a real person speaking
 
+Text to rewrite: {text}
 
-def build_prompt(text: str, style: str) -> str:
-    """Build prompt for the selected style."""
-    system = STYLE_PROMPTS[style]
-    return (
-        f"<|system|>\n{system}\n\n"
-        f"<|user|>\n{text}\n\n"
-        f"<|assistant|>\n"
-    )
+Rewritten (casual):""",
 
+    "academic": """You are an academic writing assistant. Rewrite the text below in formal scholarly language.
 
-# ---------------- MODEL CALL ---------------- #
+IMPORTANT RULES:
+1. Output ONLY the rewritten text
+2. No explanations, no notes
+3. Keep the same meaning
+4. Use precise academic vocabulary
+5. Maintain formal structure
 
-async def call_llm(prompt: str, temperature: float = 0.25):
-    loop = asyncio.get_event_loop()
+Text to rewrite: {text}
 
-    def sync_call():
-        return llm(
-            prompt,
-            max_tokens=512,
-            temperature=temperature,
-            top_p=0.9,
-            top_k=40,
-            repeat_penalty=1.1,
-            stop=[END_TOKEN],
-            echo=False,
-        )
+Rewritten (academic):""",
 
-    out = await loop.run_in_executor(None, sync_call)
+    "marketing": """You are a marketing copywriter. Rewrite the text below into persuasive marketing language.
 
-    if "choices" in out:
-        text = out["choices"][0].get("text", "")
-    else:
-        text = str(out)
+IMPORTANT RULES:
+1. Output ONLY the rewritten text
+2. No explanations, no notes
+3. Keep the same meaning
+4. Use emotional hooks and benefits
+5. Make it engaging and compelling
 
-    return clean_output(text)
+Text to rewrite: {text}
 
+Rewritten (marketing):"""
+}
 
-# ---------------- ENDPOINT: /api/transform ---------------- #
+STYLE_TEMPERATURES = {
+    "professional": 0.3,
+    "casual": 0.6,
+    "academic": 0.4,
+    "marketing": 0.7
+}
 
-@app.post("/api/transform")
-async def transform(req: TransformRequest):
-    text = req.text.strip()
+# ==================== HELPER FUNCTIONS ====================
+def load_model():
+    """Load the GGUF model"""
+    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
+                MODEL_PATH = hf_hub_download(
+                    repo_id=MODEL_REPO,
+                    filename=MODEL_FILE,
+                    local_dir=".",
+                    token=os.environ.get("HF_TOKEN", None)
+                )
+            except ImportError:
+                print("⚠️  huggingface-hub not installed, using local 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}")
+        return None
+
+def clean_output(text: str) -> str:
+    """Clean model output"""
     if not text:
-        raise HTTPException(400, "Text cannot be empty")
-
-    if req.style not in STYLE_PROMPTS:
-        raise HTTPException(400, "Invalid style")
-
-    # Marketing = more creativity
-    temperature = 0.65 if req.style == "marketing" else 0.25
-
-    prompt = build_prompt(text, req.style)
-    transformed = await call_llm(prompt, temperature)
-
+        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 whitespace
+    clean = re.sub(r'\n+', ' ', clean)
+    clean = re.sub(r'\s+', ' ', clean)
+    clean = clean.strip()
+    
+    # Remove quotes if entire text is quoted
+    if clean.startswith('"') and clean.endswith('"'):
+        clean = clean[1:-1]
+    
+    return clean
+
+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
+
+async def transform_with_model(text: str, style: str) -> str:
+    """Transform text using the loaded model"""
+    global llm
+    
+    if llm is None:
+        llm = load_model()
+        if llm is None:
+            raise HTTPException(status_code=503, detail="Model not available")
+    
+    try:
+        # Build prompt
+        prompt = format_prompt(text, style)
+        temperature = STYLE_TEMPERATURES[style]
+        
+        # Run inference
+        start_time = time.time()
+        
+        output = llm(
+            prompt,
+            max_tokens=min(400, len(text) + 100),  # Dynamic token limit
+            temperature=temperature,
+            top_p=0.9,
+            repeat_penalty=1.1,
+            stop=[END_TOKEN, "<|end|>", "\n\n"],
+            echo=False,
+        )
+        
+        processing_time = time.time() - start_time
+        
+        # Extract result
+        if "choices" in output and len(output["choices"]) > 0:
+            result = output["choices"][0]["text"]
+        else:
+            result = str(output)
+        
+        # Clean result
+        cleaned = clean_output(result)
+        
+        # Fallback if output is empty
+        if not cleaned or cleaned.isspace():
+            cleaned = f"[{style.capitalize()} Version]: {text}"
+        
+        print(f"✅ Transformation completed in {processing_time:.2f}s")
+        return cleaned
+        
+    except Exception as e:
+        print(f"❌ Model error: {e}")
+        raise HTTPException(status_code=500, detail=f"Model error: {str(e)}")
+
+# ==================== API ENDPOINTS ====================
+@app.get("/")
+async def root():
+    """Health check endpoint"""
     return {
-        "original": text,
-        "transformed": transformed,
-        "style": req.style
+        "status": "online",
+        "service": "FormatAI Humanizer",
+        "model": "Phi-3.1-mini-4k-instruct-GGUF",
+        "styles_available": list(STYLE_PROMPTS.keys()),
+        "max_input_length": MAX_INPUT_LENGTH
     }
 
-
-# ---------------- LEGACY ENDPOINT: /api/humanize ---------------- #
+@app.get("/health")
+async def health_check():
+    """Detailed health check"""
+    return {
+        "status": "healthy" if llm else "model_loading",
+        "model_loaded": llm is not None,
+        "threads": N_THREADS,
+        "context_size": N_CTX
+    }
 
 @app.post("/api/humanize")
-async def humanize(req: HumanizeRequest):
-    """Fallback endpoint — always uses casual."""
-    prompt = build_prompt(req.text.strip(), "casual")
-    out = await call_llm(prompt, temperature=0.4)
-    return {"result": out}
+async def humanize_text(request: HumanizeRequest):
+    """
+    Legacy endpoint for backward compatibility
+    Uses casual style by default
+    """
+    start_time = time.time()
+    
+    try:
+        result = await transform_with_model(request.text, "casual")
+        
+        return {
+            "result": result,
+            "original": request.text,
+            "style": "casual",
+            "processing_time": time.time() - start_time
+        }
+        
+    except HTTPException:
+        raise
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
 
+@app.post("/api/transform")
+async def transform_text(request: TransformRequest):
+    """
+    Main transformation endpoint
+    """
+    start_time = time.time()
+    
+    # Validate input
+    if not request.text or not request.text.strip():
+        raise HTTPException(status_code=400, detail="Text cannot be empty")
+    
+    if len(request.text) > MAX_INPUT_LENGTH:
+        raise HTTPException(
+            status_code=400, 
+            detail=f"Text too long. Max {MAX_INPUT_LENGTH} characters"
+        )
+    
+    if request.style not in STYLE_PROMPTS:
+        raise HTTPException(
+            status_code=400,
+            detail=f"Invalid style. Choose from: {list(STYLE_PROMPTS.keys())}"
+        )
+    
+    try:
+        # Transform the text
+        transformed = await transform_with_model(request.text, request.style)
+        
+        return {
+            "original": request.text,
+            "transformed": transformed,
+            "style": request.style,
+            "processing_time": round(time.time() - start_time, 2),
+            "success": True
+        }
+        
+    except HTTPException:
+        raise
+    except Exception as e:
+        print(f"❌ Transformation error: {e}")
+        raise HTTPException(status_code=500, detail=f"Transformation failed: {str(e)}")
+
+@app.get("/api/styles")
+async def get_styles():
+    """Get available transformation 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(".", "")
+        
+        styles_info[style] = {
+            "description": description,
+            "temperature": STYLE_TEMPERATURES[style],
+            "example_prompt": prompt[:100] + "..." if len(prompt) > 100 else prompt
+        }
+    
+    return {
+        "available_styles": styles_info,
+        "default_style": "casual"
+    }
 
-# ---------------- HEALTH CHECK ---------------- #
+# Error handler
+@app.exception_handler(Exception)
+async def general_exception_handler(request, exc):
+    return JSONResponse(
+        status_code=500,
+        content={
+            "error": "Internal server error",
+            "message": str(exc),
+            "path": request.url.path
+        }
+    )
 
-@app.get("/")
-def health():
-    return {"status": "ok", "model": MODEL_PATH}
+# ==================== MAIN ====================
+if __name__ == "__main__":
+    import uvicorn
+    
+    port = int(os.environ.get("PORT", 8000))
+    
+    uvicorn.run(
+        "app:app",
+        host="0.0.0.0",
+        port=port,
+        reload=False  # Disable reload for production
+    )