Update main.py

main.py

@@ -1,9 +1,8 @@
+import torch
 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__)
@@ -19,8 +18,6 @@ llm = LlamaCpp(
     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')
@@ -33,9 +30,9 @@ def get_skills():
 
     # Simple inference example
     output = llm(
-        f"<|user|>\n{cv_body}<|end|>\n<|assistant|>Can you list the skills mentioned in the CV?<|end|>",
+        f"\n{cv_body}\nCan you list the skills mentioned in the CV?",
         max_tokens=256,  # Generate up to 256 tokens
-        stop=["<|end|>"], 
+        stop=[""], 
         echo=True,  # Whether to echo the prompt
     )
 
@@ -47,16 +44,16 @@ def health():
 
 @app.route('/compare', methods=['POST'])
 def compare():
-    jobs_skill = request.json.get('job_skills')
+    jobs_skills = request.json.get('job_skills')
     employee_skills = request.json.get('employee_skills')
 
     # Validation
     if not isinstance(jobs_skills, list) or not all(isinstance(skill, str) for skill in jobs_skills):
-        raise ValueError("jobs_skills must be a list of strings")
+        raise ValueError("job_skills must be a list of strings")
 
     # Encoding skills into embeddings
     employee_embeddings = model.encode(employee_skills)
-   job_embeddings = model.encode(job_skills)
+    job_embeddings = model.encode(jobs_skills)
 
     # Computing cosine similarity between employee skills and each job
     similarity_scores = []
@@ -64,9 +61,10 @@ def compare():
 
     for i, job_e in enumerate(job_embeddings):
         job_e_tensor = torch.from_numpy(job_e).unsqueeze(0)
-        similarity_score = cosine_similarity(employee_embeddings_tensor, job_e_tensor, dim=1)
+        similarity_score = torch.nn.functional.cosine_similarity(employee_embeddings_tensor, job_e_tensor, dim=1)
         similarity_scores.append({"job": jobs_skills[i], "similarity_score": similarity_score.item()})
 
     return jsonify(similarity_scores)
+
 if __name__ == '__main__':
-    app.run()
+    app.run()