Create app.py

app.py

@@ -0,0 +1,55 @@
+from transformers import AutoTokenizer, AutoModel
+from datasets import load_dataset
+import torch
+
+model_ckpt = "BAAI/bge-large-en-v1.5"
+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-gpe")
+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
+
+  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=find_in_docs, inputs="text", outputs="text")
+
+demo.launch(share=True)