Add application file

Dockerfile

@@ -0,0 +1,14 @@
+FROM python:3.10
+
+RUN useradd -m -u 1000 user
+USER user
+ENV PATH="/home/user/.local/bin:$PATH"
+
+WORKDIR /app
+
+COPY --chown=user ./requirements.txt requirements.txt
+RUN pip install --no-cache-dir --upgrade -r requirements.txt
+
+COPY --chown=user . /app
+
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]

app.py

@@ -0,0 +1,71 @@
+from fastapi import FastAPI
+from pydantic import BaseModel
+from transformers import AutoTokenizer, T5ForConditionalGeneration, AutoConfig
+import torch
+
+app = FastAPI(title="CodeT5+ Backend on HuggingFace")
+
+# ==== LOAD MODEL ====
+base_ckpt = "Salesforce/codet5p-770m"
+finetuned_ckpt = "OSS-forge/codet5p-770m-pyresbugs"
+
+print("Loading tokenizer + config...")
+tokenizer = AutoTokenizer.from_pretrained(base_ckpt)
+config = AutoConfig.from_pretrained(base_ckpt)
+
+print("Loading fine-tuned model weights...")
+model = T5ForConditionalGeneration.from_pretrained(
+    finetuned_ckpt,
+    config=config
+)
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+print("Running on:", device)
+
+model = model.to(device)
+model.eval()
+
+# ==== REQUEST / RESPONSE MODELS ====
+class GenerateRequest(BaseModel):
+    prompt: str
+    language: str | None = "Python"
+    task: str = "generate"
+    max_new_tokens: int = 128
+    num_beams: int = 4
+    temperature: float = 0.7
+
+class GenerateResponse(BaseModel):
+    output: str
+
+
+def build_prompt(req: GenerateRequest):
+    if req.task == "generate":
+        return f"Generate {req.language} code:\n{req.prompt}"
+    elif req.task == "fix":
+        return f"Fix the bug in the following {req.language} code:\n{req.prompt}\n\nCorrected code:"
+    else:
+        return req.prompt
+
+
+@app.post("/generate", response_model=GenerateResponse)
+def generate(req: GenerateRequest):
+
+    prompt = build_prompt(req)
+    inputs = tokenizer(prompt, return_tensors="pt").to(device)
+
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=req.max_new_tokens,
+            num_beams=req.num_beams,
+            temperature=req.temperature,
+            early_stopping=True
+        )
+
+    text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return GenerateResponse(output=text)
+
+
+@app.get("/")
+def root():
+    return {"status": "CodeT5+ backend is running 🚀"}