Spaces:
Sleeping
Sleeping
| from fastapi import FastAPI | |
| from pydantic import BaseModel | |
| from typing import List, Dict | |
| from transformers import AutoTokenizer, AutoModel | |
| import torch | |
| app = FastAPI() | |
| # Load transformer model & tokenizer | |
| MODEL_NAME = "Worko71/keywords_new_model_without_pool" #"Worko71/ketword_match" # Can be replaced with any HuggingFace model | |
| tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME) | |
| model = AutoModel.from_pretrained(MODEL_NAME) | |
| # Utility: Mean Pooling to get sentence-level embeddings | |
| def mean_pooling(model_output, attention_mask): | |
| token_embeddings = model_output[0] # First element contains token embeddings | |
| input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float() | |
| return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp( | |
| input_mask_expanded.sum(1), min=1e-9 | |
| ) | |
| # ---------- Request Models ---------- | |
| class TextPair(BaseModel): | |
| text1: str | |
| text2: str | |
| class SkillSet(BaseModel): | |
| coreSkills: List[str] | |
| softSkills: List[str] | |
| class JobMap(BaseModel): | |
| jobs: Dict[str, SkillSet] | |
| userSkills: SkillSet | |
| class SingleJobComparison(BaseModel): | |
| userSkills: SkillSet | |
| jobDetails: SkillSet | |
| # ---------- Endpoint 1: Compare Two Sentences ---------- | |
| async def compute_similarity(payload: TextPair): | |
| """ | |
| Compare two texts and return cosine similarity score. | |
| """ | |
| inputs = tokenizer([payload.text1, payload.text2], padding=True, truncation=True, return_tensors="pt") | |
| with torch.no_grad(): | |
| model_output = model(**inputs) | |
| embeddings = mean_pooling(model_output, inputs['attention_mask']) | |
| sim_score = torch.nn.functional.cosine_similarity(embeddings[0], embeddings[1], dim=0).item() | |
| return {"similarity_percent": round(sim_score * 100, 2)} | |
| # ---------- Endpoint 2: Compare User vs Job (Soft + Core) ---------- | |
| async def compare_user_job(payload: SingleJobComparison): | |
| """ | |
| Compares a user's soft and core skills with a single job's soft and core skills. | |
| Returns individual similarity scores. | |
| """ | |
| def encode_average(texts: List[str]) -> torch.Tensor: | |
| inputs = tokenizer(texts, padding=True, truncation=True, return_tensors="pt") | |
| with torch.no_grad(): | |
| model_output = model(**inputs) | |
| embeddings = mean_pooling(model_output, inputs['attention_mask']) | |
| return embeddings.mean(dim=0) # Average across all skill embeddings | |
| user_core = encode_average(payload.userSkills.coreSkills) | |
| user_soft = encode_average(payload.userSkills.softSkills) | |
| job_core = encode_average(payload.jobDetails.coreSkills) | |
| job_soft = encode_average(payload.jobDetails.softSkills) | |
| core_score = torch.nn.functional.cosine_similarity(user_core, job_core, dim=0).item() | |
| soft_score = torch.nn.functional.cosine_similarity(user_soft, job_soft, dim=0).item() | |
| return { | |
| "similarity_percent": { | |
| "core": round(core_score * 100, 2), | |
| "soft": round(soft_score * 100, 2) | |
| } | |
| } | |
| # ---------- Endpoint 3: Bulk Comparison (User vs All Jobs) ---------- | |
| async def compare_bulk(payload: JobMap): | |
| """ | |
| Compare a user's skill set against multiple jobs. | |
| Returns a similarity map per job ID with core and soft skill scores. | |
| """ | |
| def encode_average(texts: List[str]) -> torch.Tensor: | |
| inputs = tokenizer(texts, padding=True, truncation=True, return_tensors="pt") | |
| with torch.no_grad(): | |
| model_output = model(**inputs) | |
| embeddings = mean_pooling(model_output, inputs['attention_mask']) | |
| return embeddings.mean(dim=0) | |
| user_core = encode_average(payload.userSkills.coreSkills) | |
| user_soft = encode_average(payload.userSkills.softSkills) | |
| result = {} | |
| for job_id, job in payload.jobs.items(): | |
| job_core = encode_average(job.coreSkills) | |
| job_soft = encode_average(job.softSkills) | |
| core_score = torch.nn.functional.cosine_similarity(user_core, job_core, dim=0).item() | |
| soft_score = torch.nn.functional.cosine_similarity(user_soft, job_soft, dim=0).item() | |
| result[job_id] = { | |
| "core": round(core_score * 100, 2), | |
| "soft": round(soft_score * 100, 2) | |
| } | |
| return result | |