Update app.py

app.py

@@ -1,10 +1,55 @@
-import gradio as gr
+from transformers import AutoTokenizer, AutoModel
+from datasets import load_dataset
+import torch
+
+model_ckpt = "sentence-transformers/multi-qa-mpnet-base-dot-v1"
+tokenizer = AutoTokenizer.from_pretrained(model_ckpt)
+model = AutoModel.from_pretrained(model_ckpt)
+
+device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+model.to(device)
+
+def cls_pooling(model_output):
+    return model_output.last_hidden_state[:, 0]
+
+def get_embeddings(text_list):
+    encoded_input = tokenizer(
+        text_list, padding=True, truncation=True, return_tensors="pt"
+    )
+    encoded_input = {k: v.to(device) for k, v in encoded_input.items()}
+    model_output = model(**encoded_input)
+    return cls_pooling(model_output)
+
+
+embeddings_doc_dataset = load_dataset("fashxp/pimcore-docs-embeddings")
+embeddings_doc_dataset = embeddings_doc_dataset['train']
+embeddings_doc_dataset.add_faiss_index(column="embeddings")
+
+import pandas as pd
 
+def find_in_docs(question):
+  question_embedding = get_embeddings([question]).cpu().detach().numpy()
+  question_embedding.shape
 
-def greet(name):
-    return "Hello " + name
+  scores, samples = embeddings_doc_dataset.get_nearest_examples(
+      "embeddings", question_embedding, k=10
+  )
 
+  samples_df = pd.DataFrame.from_dict(samples)
+  samples_df["scores"] = scores
+  samples_df.sort_values("scores", ascending=False, inplace=True)
+
+  result = ''
+
+  for _, row in samples_df.iterrows():
+      result = result + f"HEADING: {row.heading}\n" + f"SCORE: {row.scores}\n" + f"URL: {row.url}\n" + ("=" * 50) + "\n\n"
+
+  return result
+
+
+
+import gradio as gr
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
+demo = gr.Interface(fn=find_in_docs, inputs="text", outputs="text")
 
-demo.launch()
+demo.launch(share=True)