Update app.py

app.py

@@ -1,6 +1,4 @@
-# ================================
 #  GLOBAL WARNING SUPPRESSION
-# ================================
 import warnings
 warnings.filterwarnings("ignore", category=UserWarning)
 warnings.filterwarnings("ignore", category=FutureWarning)
@@ -8,34 +6,30 @@ warnings.filterwarnings("ignore", category=RuntimeWarning)
 warnings.filterwarnings("ignore", category=DeprecationWarning)
 warnings.simplefilter("ignore")
 
-# ================================
 #  IMPORTS
-# ================================
 import json
 import pickle
 import numpy as np
 import torch
+import io 
 from io import StringIO
 from Bio import SeqIO
+import csv
 
 from fastapi import FastAPI, Request, UploadFile, File
 from fastapi.staticfiles import StaticFiles
-from fastapi.responses import HTMLResponse
+from fastapi.responses import HTMLResponse StreamingResponse
 from fastapi.templating import Jinja2Templates
 
 from transformers import AutoTokenizer, AutoModel
 
-# ================================
 #  FASTAPI INIT + MOUNTS
-# ================================
 app = FastAPI()
 
 app.mount("/static", StaticFiles(directory="static"), name="static")
 templates = Jinja2Templates(directory="templates")
 
-# ================================
 #  LOAD MODEL + TOKENIZER
-# ================================
 DEVICE = torch.device("cpu")
 
 tokenizer = AutoTokenizer.from_pretrained("facebook/esm2_t30_150M_UR50D")
@@ -50,9 +44,8 @@ with open("label_map.json", "r") as f:
 
 INV_LABEL_MAP = {v: k for k, v in LABEL_MAP.items()}
 
-# ================================
 #  ESM2 EMBEDDING FUNCTION
-# ================================
+
 def embed_sequence(seq: str) -> np.ndarray:
   seq = seq.strip()
   inputs = tokenizer(seq, return_tensors="pt", add_special_tokens=True)
@@ -65,9 +58,8 @@ def embed_sequence(seq: str) -> np.ndarray:
   mean_emb = token_emb[1:-1].mean(dim=0)
   return mean_emb.cpu().numpy().reshape(1, -1)
 
-# ================================
 #  PREDICT ONE SEQUENCE
-# ================================
+
 def run_single_prediction(seq: str):
   emb = embed_sequence(seq)
   probs = classifier.predict_proba(emb)[0]
@@ -79,9 +71,8 @@ def run_single_prediction(seq: str):
     "probabilities": {INV_LABEL_MAP[i]: float(p) for i, p in enumerate(probs)}
   }
 
-# ================================
 #  PREDICT FASTA FILE
-# ================================
+
 def run_fasta_prediction(content: str):
   results = []
   handle = StringIO(content)
@@ -105,9 +96,8 @@ def run_fasta_prediction(content: str):
 
   return {"results": results}
 
-# ================================
 #  PAGE ROUTES
-# ================================
+
 @app.get("/", response_class=HTMLResponse)
 async def home(request: Request):
   return templates.TemplateResponse("index.html", {"request": request})
@@ -124,9 +114,9 @@ async def help_page(request: Request):
 async def contact(request: Request):
   return templates.TemplateResponse("contact.html", {"request": request})
 
-# ================================
+
 #  API: UNIVERSAL SEQUENCE PREDICTION
-# ================================
+
 @app.post("/api/predict_sequence")
 async def api_predict_sequence(request: Request):
   # 1. Try JSON
@@ -147,11 +137,27 @@ async def api_predict_sequence(request: Request):
 
   return {"error": "No sequence provided"}
 
-# ================================
+
 #  API: FASTA PREDICTION
-# ================================
+
 @app.post("/api/predict_fasta")
 async def api_predict_fasta(file: UploadFile = File(...)):
   raw = await file.read()
   content = raw.decode("utf-8", errors="ignore")
   return run_fasta_prediction(content)
+
+# DOWNLOAD RESULT
+
+@app.post("/api/download_csv")
+async def download_csv(results: list[dict]):
+    output = io.StringIO()
+    writer = csv.DictWriter(output, fieldnames=results[0].keys())
+    writer.writeheader()
+    writer.writerows(results)
+    output.seek(0)
+
+    return StreamingResponse(
+        output,
+        media_type="text/csv",
+        headers={"Content-Disposition": "attachment; filename=canloc_results.csv"}
+    )