Refactor get_closest_budget_item function to accept a parameter for the number of results

app.py

@@ -28,10 +28,10 @@ budget_items = df.apply(get_name, axis=1).unique().tolist()
 
 budget_item_embeddings = torch.stack(list(map(get_embedding, budget_items)))
 
-def get_closest_budget_item(text):
+def get_closest_budget_item(text, num_results=5):
     text_embedding = get_embedding(text)
-    scores = (budget_item_embeddings * text_embedding).sum(axis=1)
-    top_idx = scores.argsort(descending=True)[:5]
+    scores = torch.norm(budget_item_embeddings - text_embedding, dim=1)
+    top_idx = scores.argsort()[:num_results]
     return pd.DataFrame({
         'budget_item': np.array(budget_items)[top_idx],
         'score': scores[top_idx].tolist()
@@ -39,7 +39,7 @@ def get_closest_budget_item(text):
 
 demo = gr.Interface(
     fn=get_closest_budget_item,
-    inputs='textbox',
+    inputs=['textbox', gr.Slider(minimum=1, maximum=50, step=5, default=5, label="Number of results")],
     outputs='dataframe',
 )