Update main.py

main.py

@@ -19,21 +19,38 @@ 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')
 print("model size ====> :", file_size.st_size, "bytes")
 
 
+@app.route('/cv', methods=['POST'])
+def get_skills():
+    cv_body = request.json.get('cv_body')
 
-# Assume 'model' is already defined somewhere else
+    # 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', methods=['POST'])
 @app.route('/compare', methods=['POST'])
 def compare():
-    employee_skills = request.json.get('jobs_skills')  
-    jobs_skills = request.json.get('employee_skills')
+    employee_skills = request.json.get('jobs_skills') 
+    jobs_skills = request.json.get('employee_skills')  
 
+    
 
     # Validation
     if not isinstance(jobs_skills, list) or not all(isinstance(skill, str) for skill in jobs_skills):
@@ -49,9 +66,11 @@ 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)[0]
+        similarity_score = 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()