Add application file

.dockerignore

@@ -0,0 +1 @@
+__pycache__

Dockerfile

@@ -0,0 +1,17 @@
+# Base Image
+FROM python:3.12
+
+# Ensure the Python output is not buffered
+ENV PYTHONUNBUFFERED True
+
+# Set the working directory in the container
+WORKDIR /app
+
+# Copy all the application files to the working directory
+COPY . /app
+
+# Install the required Python packages
+RUN pip install -r requirements.txt
+
+# Command to run the web service
+CMD ["python", "main.py"]

main.py

@@ -0,0 +1,35 @@
+import re
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+
+model_dir = 'edithram23/Redaction'
+tokenizer = AutoTokenizer.from_pretrained(model_dir)
+model = AutoModelForSeq2SeqLM.from_pretrained(model_dir)
+
+def mask_generation(text):
+    import re
+    inputs = ["Mask Generation: " + text]
+    inputs = tokenizer(inputs, max_length=500, truncation=True, return_tensors="pt")
+    output = model.generate(**inputs, num_beams=8, do_sample=True, max_length=len(text)+10)
+    decoded_output = tokenizer.batch_decode(output, skip_special_tokens=True)[0]
+    predicted_title = decoded_output.strip()
+    pattern = r'\[.*?\]'
+    # Replace all occurrences of the pattern with [redacted]
+    redacted_text = re.sub(pattern, '[redacted]', predicted_title)
+    return redacted_text
+
+from fastapi import FastAPI
+import uvicorn
+
+app = FastAPI()
+
+@app.get("/")
+async def hello():
+    return {"msg" : "Live"}
+
+@app.post("/mask")
+async def mask_input(query):
+    output = mask_generation(query)
+    return {"data" : output}
+
+if __name__ == '__main__':
+    uvicorn.run("main:app", host="0.0.0.0", port=7860, reload=True, workers=1)

requirements.txt

@@ -0,0 +1,4 @@
+fastapi==0.111.0
+transformers==4.41.2
+uvicorn==0.30.1
+torch