omarz / main.py
zayanomar5's picture
Update main.py
f934933 verified
raw
history blame
2.23 kB
from flask import Flask, request, jsonify
from langchain_community.llms import LlamaCpp
from sentence_transformers import SentenceTransformer
from transformers import AutoModel
import torch
from torch.nn.functional import cosine_similarity
import os
app = Flask(__name__)
n_gpu_layers = 0
n_batch = 1024
llm = LlamaCpp(
model_path="Phi-3-mini-4k-instruct-q4.gguf", # path to GGUF file
temperature=0.1,
n_gpu_layers=n_gpu_layers,
n_batch=n_batch,
verbose=True,
n_ctx=4096
)
model0 = AutoModel.from_pretrained('sentence-transformers/paraphrase-TinyBERT-L6-v2')
model = SentenceTransformer('sentence-transformers/paraphrase-TinyBERT-L6-v2')
file_size = os.stat('Phi-3-mini-4k-instruct-q4.gguf')
print("model size ====> :", file_size.st_size, "bytes")
@app.route('/cv', methods=['POST'])
def get_skills():
cv_body = request.json.get('cv_body')
# Simple inference example
output = llm(
f"<|user|>\n{cv_body}<|end|>\n<|assistant|>Can you list the skills mentioned in the CV?<|end|>",
max_tokens=256, # Generate up to 256 tokens
stop=["<|end|>"],
echo=True, # Whether to echo the prompt
)
return jsonify({'skills': output})
@app.get('/')
def health():
return jsonify({'status': 'Worked'})
@app.route('/compare_job', methods=['POST'])
def compare_job():
job_skills = request.json.get('post')
employee_skills = request.json.get('job_skills')
# Validation
if not isinstance(employee_skills, list) or not all(isinstance(skill, str) for skill in employee_skills):
raise ValueError("employee_skills must be a list of strings")
# Encoding skills into embeddings
job_embedding = model.encode(job_skills)
employee_embedding = model.encode(employee_skills)
# Computing cosine similarity between employee skills and the job
employee_embedding_tensor = torch.from_numpy(employee_embedding).unsqueeze(0)
job_embedding_tensor = torch.from_numpy(job_embedding).unsqueeze(0)
similarity_score = cosine_similarity(employee_embedding_tensor, job_embedding_tensor, dim=1)
return jsonify({"job": job_skills, "similarity_score": similarity_score.item()})
if __name__ == '__main__':
app.run()