Views
No views yet
👉For full documentation, see: https://github.com/GeoGPT-Research-Project/GeoGPT-RAG
1import numpy as np
2from sentence_transformers import SentenceTransformer
3
4task_description = 'Given a web search query, retrieve relevant passages that answer the query'
5def get_detailed_instruct(task_description: str, query: str) -> str:
6 return f'Instruct: {task_description}\nQuery: {query}'
7
8model_name_or_path = 'GeoGPT/GeoEmbedding'
9
10model = SentenceTransformer(model_name_or_path, device="cuda", trust_remote_code=True)
11
12queries = [
13 "What is the main cause of earthquakes?",
14 "How do sedimentary rocks form?",
15]
16
17passages = [
18 "Earthquakes occur due to the sudden release of energy in the Earth's crust, often caused by tectonic plate movements along fault lines.",
19 "Sedimentary rocks form through the deposition and compaction of mineral and organic particles over time, typically in water bodies.",
20]
21
22queries = [get_detailed_instruct(task_description, query) for query in queries]
23
24q_vecs = model.encode(queries, normalize_embeddings=True)
25p_vecs = model.encode(passages, normalize_embeddings=True)
26
27print(np.dot(q_vecs, p_vecs.T))
28#[[0.6369 0.2092 ]
29# [0.2499 0.8411 ]]