ticketsAPI / app.py
Tonycrux's picture
Update app.py
b39c824 verified
Raw History Blame Contribute Delete
2.37 kB
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
@app.get("/")
def root():
return {"status": "running", "message": "Use POST /classify"}
@app.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}
@app.post("/similarity")
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}
@app.post("/gender")
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)