Update app.py

app.py

@@ -32,26 +32,19 @@ index.add(embeddings)
 
 # Search for similar documents
 query = "训练脚本."
-input_ids = tokenizer.encode(query, return_tensors="pt")
-with torch.no_grad():
-    query_embedding = model(input_ids)[0][0].numpy()
-k = 2  # Number of similar documents to retrieve
-D, I = index.search(query_embedding.reshape(1, -1), k)
 
-# Print the results
-st.write(f"Query: {query}")
-for i in range(k):
-    st.write(f"Rank {i+1}: {texts[I[0][i]]} (similarity score: {D[0][i]})")
+from sklearn.metrics.pairwise import cosine_similarity
 
-    
-# Search index for the most similar content
-k = 5  # Number of results to retrieve
-D, I = index.search(np.array([question_embedding]), k)
-
-# Display the results
-st.write("Top {} similar content:".format(k))
-for i in range(k):
-    st.write("{}: {} : {}".format(i+1, knowledge[I[0][i]], I[0][i]))
+query_embedding = model.encode([query])[0]
+
+# Compute the cosine similarity between the query embedding and the document embeddings
+similarities = cosine_similarity([query_embedding], embeddings)[0]
+
+# Get the index of the most similar document
+most_similar_index = similarities.argmax()
+
+# Print the most similar document
+st.write(documents[most_similar_index])
 
 st.title('My first app')