Update backend/inference.py

backend/inference.py

@@ -1,5 +1,7 @@
 import torch
-from backend.utils import load_model, load_embeddings, load_texts
+
+from backend.utils import load_embeddings, load_model, load_texts
+
 
 # Search
 def query_search(query: str, n_answers: int, model_name: str):
@@ -20,6 +22,8 @@ def query_search(query: str, n_answers: int, model_name: str):
 
     corpus_texts["Similarity"] = hits.tolist()
 
+    print(corpus_texts)
+
     return corpus_texts.sort_values(by="Similarity", ascending=False).head(n_answers)[
-        ["Description", "Code"]
+        ["func_documentation_string", "repository_name", "func_code_url"]
     ]