ticketsAPI / app.py
Tonycrux's picture
Add similarity endpoint
40f0642 verified
raw
history blame
1.49 kB
import re
import joblib
from fastapi import FastAPI, Request
from pydantic import BaseModel
from fastapi.middleware.cors import CORSMiddleware.
from sentence_transformers import SentenceTransformer, util
app = FastAPI()
app.add_middleware(
CORSMiddleware,
allow_origins=["*"], # Or restrict to your domain
allow_methods=["*"],
allow_headers=["*"],
)
# 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")
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}