added cache

Dockerfile

@@ -1,10 +1,33 @@
+# Use an official lightweight Python base image
 FROM python:3.10-slim
 
+# Set environment variables
+ENV PYTHONDONTWRITEBYTECODE=1
+ENV PYTHONUNBUFFERED=1
+ENV TRANSFORMERS_CACHE=/data
+
+# Set working directory
 WORKDIR /app
 
+# Install system dependencies
+RUN apt-get update && apt-get install -y \
+    git \
+    && rm -rf /var/lib/apt/lists/*
+
+# Copy requirements (you must have a requirements.txt)
 COPY requirements.txt .
+
+# Install Python dependencies
 RUN pip install --no-cache-dir -r requirements.txt
 
+# Copy entire project
 COPY . .
 
-CMD ["uvicorn", "app.main:app", "--host", "0.0.0.0", "--port", "7860"]
+# Create /data dir for model caching
+RUN mkdir -p /data
+
+# Expose port (optional in HF Spaces)
+EXPOSE 7860
+
+# Run FastAPI with Uvicorn
+CMD ["uvicorn", "app.main:app", "--host", "0.0.0.0", "--port", "7860"]

app/main.py

@@ -1,6 +1,11 @@
 from fastapi import FastAPI, Body
 from pydantic import BaseModel
 from app import models, prompts
+import logging
+
+# Setup basic logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
 app = FastAPI()
 
@@ -10,36 +15,61 @@ class TextInput(BaseModel):
     tone: str = None
     target_lang: str = None
 
+@app.get("/")
+def root():
+    return {"message": "Welcome to Build Your Own Grammarly API. Use /rewrite POST endpoint."}
+
 @app.post("/rewrite")
 def rewrite(input: TextInput):
     text = input.text
     mode = input.mode
+    logger.info(f"Received request | Mode: {mode} | Text: {text}")
+
+    try:
+        if mode == "grammar":
+            result = models.run_grammar_correction(text)
+            logger.info(f"Grammar corrected: {result}")
+            return {"result": result}
 
-    if mode == "grammar":
-        return {"result": models.run_grammar_correction(text)}
+        elif mode == "paraphrase":
+            prompt = prompts.paraphrase_prompt(text)
+            result = models.run_flan_prompt(prompt)
+            logger.info(f"Paraphrased: {result}")
+            return {"result": result}
 
-    elif mode == "paraphrase":
-        prompt = prompts.paraphrase_prompt(text)
-        return {"result": models.run_flan_prompt(prompt)}
+        elif mode == "clarity":
+            prompt = prompts.clarity_prompt(text)
+            result = models.run_flan_prompt(prompt)
+            logger.info(f"Clarity enhanced: {result}")
+            return {"result": result}
 
-    elif mode == "clarity":
-        prompt = prompts.clarity_prompt(text)
-        return {"result": models.run_flan_prompt(prompt)}
+        elif mode == "fluency":
+            prompt = prompts.fluency_prompt(text)
+            result = models.run_flan_prompt(prompt)
+            logger.info(f"Fluency improved: {result}")
+            return {"result": result}
 
-    elif mode == "fluency":
-        prompt = prompts.fluency_prompt(text)
-        return {"result": models.run_flan_prompt(prompt)}
+        elif mode == "tone" and input.tone:
+            prompt = prompts.tone_prompt(text, input.tone)
+            result = models.run_flan_prompt(prompt)
+            logger.info(f"Tone changed to {input.tone}: {result}")
+            return {"result": result}
 
-    elif mode == "tone" and input.tone:
-        prompt = prompts.tone_prompt(text, input.tone)
-        return {"result": models.run_flan_prompt(prompt)}
+        elif mode == "translate" and input.target_lang:
+            result = models.run_translation(text, input.target_lang)
+            logger.info(f"Translated to {input.target_lang}: {result}")
+            return {"result": result}
 
-    elif mode == "translate" and input.target_lang:
-        return {"result": models.run_translation(text, input.target_lang)}
+        elif mode == "pronoun":
+            prompt = prompts.pronoun_friendly_prompt(text)
+            result = models.run_flan_prompt(prompt)
+            logger.info(f"Inclusive pronoun version: {result}")
+            return {"result": result}
 
-    elif mode == "pronoun":
-        prompt = prompts.pronoun_friendly_prompt(text)
-        return {"result": models.run_flan_prompt(prompt)}
+        else:
+            logger.error("Invalid request parameters.")
+            return {"error": "Invalid request"}
 
-    else:
-        return {"error": "Invalid request"}
+    except Exception as e:
+        logger.error(f"Error while processing: {str(e)}")
+        return {"error": str(e)}

app/test_main.py

@@ -0,0 +1,39 @@
+from fastapi.testclient import TestClient
+from app.main import app
+
+client = TestClient(app)
+
+def test_grammar():
+    response = client.post("/rewrite", json={"text": "She go to school yesterday.", "mode": "grammar"})
+    assert response.status_code == 200
+    assert "result" in response.json()
+
+def test_paraphrase():
+    response = client.post("/rewrite", json={"text": "I love programming.", "mode": "paraphrase"})
+    assert response.status_code == 200
+    assert "result" in response.json()
+
+def test_clarity():
+    response = client.post("/rewrite", json={"text": "This sentence is a bit confusing.", "mode": "clarity"})
+    assert response.status_code == 200
+    assert "result" in response.json()
+
+def test_fluency():
+    response = client.post("/rewrite", json={"text": "He no went there before.", "mode": "fluency"})
+    assert response.status_code == 200
+    assert "result" in response.json()
+
+def test_tone():
+    response = client.post("/rewrite", json={"text": "Leave me alone.", "mode": "tone", "tone": "friendly"})
+    assert response.status_code == 200
+    assert "result" in response.json()
+
+def test_translation():
+    response = client.post("/rewrite", json={"text": "How are you?", "mode": "translate", "target_lang": "fr"})
+    assert response.status_code == 200
+    assert "result" in response.json()
+
+def test_pronoun():
+    response = client.post("/rewrite", json={"text": "Each user must provide his email.", "mode": "pronoun"})
+    assert response.status_code == 200
+    assert "result" in response.json()