Update app.py

app.py

@@ -1,146 +1,187 @@
+# app.py
+import asyncio
+import re
+from typing import Literal
 from fastapi import FastAPI, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
 from llama_cpp import Llama
-import re
 
-# GGUF MODEL
-MODEL_REPO = "bartowski/Phi-3.1-mini-4k-instruct-GGUF"
-MODEL_FILE = "Phi-3.1-mini-4k-instruct-IQ2_M.gguf"
-
-print("Loading Phi-3.1 Mini GGUF model...")
-llm = Llama.from_pretrained(
-    repo_id=MODEL_REPO,
-    filename=MODEL_FILE,
-    n_threads=4,           # Adjust based on CPU cores
-    n_ctx=4096,            # Increased to model's full capacity
-    n_batch=512,           # Increased batch size for better performance
-    n_gpu_layers=0,        # Set to >0 if you have GPU
+# ---------------- MODEL CONFIG ---------------- #
+
+# IMPORTANT: For HuggingFace Spaces, the model file is inside the repo folder
+MODEL_PATH = "model/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
+
+# Unique token to force controlled stopping
+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,
 )
 print("Model loaded successfully.")
 
-app = FastAPI()
+# ---------------- FASTAPI APP ---------------- #
+
+app = FastAPI(title="FormatAI Humanizer Backend")
 
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],
+    allow_origins=["*"],     # allow all origins (Vercel frontend)
     allow_methods=["*"],
     allow_headers=["*"],
 )
 
-class RequestBody(BaseModel):
+# ---------------- REQUEST MODELS ---------------- #
+
+class TransformRequest(BaseModel):
     text: str
+    style: Literal["professional", "casual", "academic", "marketing"]
 
-def clean_output(text: str) -> str:
-    """Clean the model output by removing extra tags and whitespace."""
-    # Remove common ending tokens and system tags
-    clean_text = re.sub(r'\[/?(SYSTEM|USER|ASSISTANT)\]', '', text)
-    clean_text = re.sub(r'</?s>', '', clean_text)
-    clean_text = re.sub(r'\s+', ' ', clean_text)  # Normalize whitespace
-    clean_text = clean_text.strip()
-    
-    # Remove the original text if it's repeated
-    lines = clean_text.split('\n')
-    if len(lines) > 1:
-        # Take the most human-like line (usually the last one)
-        clean_text = lines[-1].strip()
-    
-    return clean_text
+class HumanizeRequest(BaseModel):  # legacy
+    text: str
 
-@app.post("/api/humanize")
-async def humanize(body: RequestBody):
-    text = body.text.strip()
-    if not text:
-        raise HTTPException(status_code=400, detail="Text cannot be empty")
-
-    # More specific system prompt for humanization
-    system_prompt = """You are an expert text humanizer. Your task is to rewrite text to sound more natural and human-like while preserving the original meaning. Follow these rules:
-    1. Make the text sound like a real person wrote it
-    2. Use casual, conversational language when appropriate
-    3. Fix robotic or overly formal phrasing
-    4. Keep technical terms when necessary but explain them naturally
-    5. Maintain the original intent and tone
-    6. Output ONLY the humanized version, no explanations
-    
-    Examples:
-    Input: "What is the process for obtaining a permit?"
-    Output: "How do I get a permit?"
-    
-    Input: "The application must be submitted by the deadline."
-    Output: "You'll need to submit your application before the deadline."
-    
-    Input: "Can you provide the requested information?"
-    Output: "Could you share that information with me?"
-    
-    Now humanize this text:"""
-    
-    # Create the prompt with proper Phi-3.1 formatting
-    user_prompt = f"{text}"
-    
-    # For Phi-3.1, use this format:
-    full_prompt = f"<|user|>\n{system_prompt}\n\n{user_prompt}<|end|>\n<|assistant|>\n"
-
-    try:
-        # Generate with more tokens and slightly higher temperature for creativity
-        output = llm(
-            full_prompt,
+
+# ---------------- 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. Output ONLY the rewritten text, "
+        f"then write {END_TOKEN}."
+    ),
+
+    "casual": (
+        "STYLE: CASUAL\n"
+        "Rewrite the user's text in a friendly, natural, conversational tone. Use contractions and human-like flow. "
+        "Do NOT add explanations. Output ONLY the rewritten text, then write "
+        f"{END_TOKEN}."
+    ),
+
+    "academic": (
+        "STYLE: ACADEMIC\n"
+        "Rewrite the user's text in formal academic language suitable for scholarly work. "
+        "Use precise and objective vocabulary. Do NOT add explanations. Output ONLY the rewritten text, "
+        f"then write {END_TOKEN}."
+    ),
+
+    "marketing": (
+        "STYLE: MARKETING\n"
+        "Rewrite the user's text in persuasive, benefit-focused marketing copy. "
+        "Use strong emotional hooks and punchy messaging. Do NOT add explanations. Output ONLY the rewritten text, "
+        f"then write {END_TOKEN}."
+    ),
+}
+
+
+# ---------------- HELPERS ---------------- #
+
+def clean_output(raw: str) -> str:
+    """Strip junk tokens and trim to final output."""
+    if not raw:
+        return ""
+
+    # Remove system markers
+    raw = re.sub(r"<\|/?(system|assistant|user|end)\|>", "", raw, flags=re.I)
+
+    # Stop at END_TOKEN
+    if END_TOKEN in raw:
+        raw = raw.split(END_TOKEN)[0]
+
+    raw = raw.strip()
+    raw = re.sub(r"[ \t]+", " ", raw)
+    return raw.strip()
+
+
+def build_prompt(text: str, style: str) -> str:
+    """Create strict prompt for selected style."""
+    system = STYLE_PROMPTS[style]
+
+    return (
+        f"<|system|>\n{system}\n\n"
+        f"<|user|>\n{text}\n\n"
+        f"<|assistant|>\n"
+    )
+
+
+# ---------------- MODEL CALL ---------------- #
+
+async def call_llm(prompt: str, temperature: float = 0.25):
+    loop = asyncio.get_event_loop()
+
+    def sync_call():
+        return llm(
+            prompt,
             max_tokens=512,
-            temperature=0.7,      # Higher for more creativity
-            top_p=0.9,           # Slightly less than 1.0 for diversity
-            repeat_penalty=1.1,  # Penalize repetition
-            top_k=40,            # Limit token choices
-            echo=False,          # Don't echo the prompt in output
-            stop=["<|end|>", "<|user|>", "<|assistant|>", "\n\n", "###"],
+            temperature=temperature,
+            top_p=0.9,
+            top_k=40,
+            repeat_penalty=1.1,
+            stop=[END_TOKEN],
+            echo=False,
         )
-        
-        result = output["choices"][0]["text"].strip()
-        cleaned_result = clean_output(result)
-        
-        # If the output is empty or too similar to input, use a fallback
-        if not cleaned_result or cleaned_result == text:
-            # Simple fallback humanization rules
-            fallback_rules = [
-                ("what is", "what's"),
-                ("could you", "can you"),
-                ("would you", "could you"),
-                ("please be advised that", ""),
-                ("it is recommended that", "you should"),
-                ("utilize", "use"),
-                ("commence", "start"),
-                ("terminate", "end"),
-                ("approximately", "about"),
-                ("therefore", "so"),
-                ("however", "but"),
-                ("in order to", "to"),
-                ("with regard to", "about"),
-                ("at this point in time", "now"),
-                ("due to the fact that", "because"),
-                ("prior to", "before"),
-                ("subsequent to", "after"),
-            ]
-            
-            cleaned_result = text
-            for formal, casual in fallback_rules:
-                if formal in cleaned_result.lower():
-                    # Replace while preserving case
-                    pattern = re.compile(re.escape(formal), re.IGNORECASE)
-                    cleaned_result = pattern.sub(casual, cleaned_result)
-            
-            # Add contractions
-            cleaned_result = re.sub(r'\b(I am|you are|he is|she is|it is|we are|they are)\b', 
-                                  lambda m: m.group(1).replace(' ', "'"), cleaned_result, flags=re.IGNORECASE)
-            cleaned_result = re.sub(r'\b(do not|does not|did not|cannot|will not|would not|could not|should not|is not|are not|was not|were not)\b',
-                                  lambda m: m.group(1).replace(' ', "'").replace("cannot", "can't"), cleaned_result, flags=re.IGNORECASE)
-        
-        return {"result": cleaned_result}
-
-    except Exception as e:
-        print(f"Model error: {e}")
-        # Return a simple humanized version as fallback
-        simple_humanized = text.replace("?", "?").replace(".", ".")
-        return {"result": simple_humanized}
+
+    out = await loop.run_in_executor(None, sync_call)
+
+    if "choices" in out:
+        text = out["choices"][0].get("text", "")
+    else:
+        text = str(out)
+
+    return clean_output(text)
+
+
+# ---------------- ENDPOINT: /api/transform ---------------- #
+
+@app.post("/api/transform")
+async def transform(req: TransformRequest):
+    text = req.text.strip()
+    if not text:
+        raise HTTPException(400, "Text cannot be empty")
+
+    if req.style not in STYLE_PROMPTS:
+        raise HTTPException(400, "Invalid style")
+
+    # More creativity for marketing
+    temperature = 0.65 if req.style == "marketing" else 0.25
+
+    prompt = build_prompt(text, req.style)
+    transformed = await call_llm(prompt, temperature=temperature)
+
+    return {
+        "original": text,
+        "transformed": transformed,
+        "style": req.style
+    }
+
+
+# ---------------- LEGACY ENDPOINT: /api/humanize ---------------- #
+
+@app.post("/api/humanize")
+async def humanize(req: HumanizeRequest):
+    """Old endpoint - always uses casual."""
+    prompt = build_prompt(req.text.strip(), "casual")
+    out = await call_llm(prompt, temperature=0.4)
+    return {"result": out}
+
+
+# ---------------- HEALTH CHECK ---------------- #
 
 @app.get("/")
 def health():
-    return {"status": "ok", "model": MODEL_FILE}
+    return {"status": "ok", "model": MODEL_PATH}