Spaces:
Running
Running
import re | |
import os | |
import joblib | |
from fastapi import FastAPI, File, UploadFile | |
from pydantic import BaseModel | |
from fastapi.middleware.cors import CORSMiddleware | |
from sentence_transformers import SentenceTransformer, util | |
import cv2 | |
import numpy as np | |
app = FastAPI() | |
app.add_middleware( | |
CORSMiddleware, | |
allow_origins=["*"], # Or restrict to your domain | |
allow_methods=["*"], | |
allow_headers=["*"], | |
) | |
os.environ["HF_HOME"] = "/tmp" | |
os.environ["TRANSFORMERS_CACHE"] = "/tmp" | |
os.environ["SENTENCE_TRANSFORMERS_HOME"] = "/tmp" | |
# Load model and vectorizer | |
model = joblib.load("team_classifier_model.joblib") | |
vectorizer = joblib.load("tfidf_vectorizer.joblib") | |
sbert_model = SentenceTransformer("sentence-transformers/paraphrase-MiniLM-L6-v2") | |
gender_list = ['Male', 'Female'] | |
model = cv2.dnn.readNetFromCaffe("gender_deploy.prototxt", "gender_net.caffemodel") | |
def clean_text(text): | |
text = re.sub(r"\s+", " ", str(text)) | |
text = re.sub(r"[^\w\s]", "", text) | |
return text.lower().strip() | |
class InputText(BaseModel): | |
subject: str | |
message: str | |
class SimilarityRequest(BaseModel): | |
text1: str | |
text2: str | |
def root(): | |
return {"status": "running", "message": "Use POST /classify"} | |
async def classify_ticket(data: InputText): | |
combined = clean_text(f"{data.subject} {data.message}") | |
vec = vectorizer.transform([combined]) | |
prediction = model.predict(vec)[0] | |
return {"team": prediction} | |
async def compute_similarity(data: SimilarityRequest): | |
emb1 = sbert_model.encode(data.text1, convert_to_tensor=True) | |
emb2 = sbert_model.encode(data.text2, convert_to_tensor=True) | |
score = util.pytorch_cos_sim(emb1, emb2).item() | |
return {"similarity": score} | |
async def predict_gender(file: UploadFile = File(...)): | |
try: | |
contents = await file.read() | |
npimg = np.frombuffer(contents, np.uint8) | |
img = cv2.imdecode(npimg, cv2.IMREAD_COLOR) | |
blob = cv2.dnn.blobFromImage(img, 1.0, (227, 227), (78.426337, 87.768914, 114.895847), swapRB=False) | |
model.setInput(blob) | |
gender_preds = model.forward() | |
gender = gender_list[gender_preds[0].argmax()] | |
return {"gender": gender} | |
except Exception as e: | |
return JSONResponse(content={"error": str(e)}, status_code=500) |