arafamustafa's picture
Update app.py
acc3bdf verified
Raw
History Blame Contribute Delete
4.08 kB
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from sentence_transformers import SentenceTransformer
import torch
import pandas as pd
import numpy as np
import os
import gdown
# -----------------------------
# Download files from Google Drive
# -----------------------------
MODEL_DRIVE_ID = "1qdW9UCJlZSSlSRahx30UrLtZ_3MRMm9e"
CSV_DRIVE_ID = "1lzJU7hjG75uQ4hL02Gzy_vU1FtiHOImT"
if not os.path.exists("trained_model"):
print("Downloading model from Google Drive...")
os.makedirs("trained_model", exist_ok=True)
gdown.download_folder(id=MODEL_DRIVE_ID, output="trained_model", quiet=False)
print("✓ Model downloaded")
if not os.path.exists("projects.csv"):
print("Downloading projects.csv from Google Drive...")
gdown.download(id=CSV_DRIVE_ID, output="projects.csv", quiet=False)
print("✓ projects.csv downloaded")
# -----------------------------
# App initialization
# -----------------------------
app = FastAPI(
title="Project Idea Similarity API",
description="Semantic similarity search for graduation projects",
version="1.0"
)
def clean_text(text: str) -> str:
text = str(text).strip()
text = " ".join(text.split())
return text
print("Loading model...")
model = SentenceTransformer("trained_model")
print("✓ Model loaded")
print("Loading dataset...")
df = pd.read_csv("projects.csv", encoding="latin1")
df.rename(columns={'Name': 'Name'}, inplace=True)
df["project_id"] = df.index + 1
required_cols = {"Name", "Abstract", "Date"}
print("Number of records:", len(df))
if not required_cols.issubset(df.columns):
raise RuntimeError("Dataset must contain Name, Abstract, Date columns")
df["Abstract"] = df["Abstract"].apply(clean_text)
projects = df.to_dict(orient="records")
EMBEDDINGS_PATH = "/tmp/embeddings.npy"
try:
print("Loading embeddings...")
embeddings = np.load(EMBEDDINGS_PATH)
embeddings = torch.tensor(embeddings)
print("✓ Embeddings loaded")
except FileNotFoundError:
print("Computing embeddings...")
embeddings = model.encode(
df["Abstract"].tolist(),
convert_to_tensor=True,
show_progress_bar=True
)
np.save(EMBEDDINGS_PATH, embeddings.cpu().numpy())
print("✓ Embeddings computed & saved")
# -----------------------------
# Models
# -----------------------------
class SimilarityRequest(BaseModel):
abstract: str
top_k: int = 5
class SimilarProject(BaseModel):
project_id: int
project_name: str
project_abstract: str
project_date: str
similarity_score: float
class SimilarityResponse(BaseModel):
input_abstract: str
top_k: int
results: list[SimilarProject]
# -----------------------------
# Routes
# -----------------------------
@app.get("/api/health")
def health():
return {
"status": "ok",
"model_loaded": True,
"projects_loaded": len(projects)
}
@app.post("/api/similarity", response_model=SimilarityResponse)
def find_similarity(request: SimilarityRequest):
if not request.abstract or len(request.abstract.strip()) < 20:
raise HTTPException(
status_code=400,
detail="Project abstract must be at least 20 characters"
)
cleaned_input = clean_text(request.abstract)
input_embedding = model.encode(cleaned_input, convert_to_tensor=True)
scores = torch.nn.functional.cosine_similarity(
input_embedding.unsqueeze(0),
embeddings
)
top_k = min(request.top_k, len(projects))
top_scores, top_indices = torch.topk(scores, k=top_k)
results = []
for score, idx in zip(top_scores, top_indices):
project = projects[idx.item()]
results.append({
"project_id": project["project_id"]+1,
"project_name": project["Name"],
"project_abstract": project["Abstract"],
"project_date": str(project["Date"]),
"similarity_score": round(score.item(), 4)
})
return {
"input_abstract": cleaned_input,
"top_k": top_k,
"results": results
}