Gurjot Singh
commited on
Commit
·
5b5f653
1
Parent(s):
d5d0856
fixed linting errors
Browse files
examples/lightrag_gemini_demo.py
CHANGED
@@ -47,20 +47,19 @@ async def llm_model_func(
|
|
47 |
response = client.models.generate_content(
|
48 |
model="gemini-1.5-flash",
|
49 |
contents=[combined_prompt],
|
50 |
-
config=types.GenerateContentConfig(
|
51 |
-
max_output_tokens=500,
|
52 |
-
temperature=0.1
|
53 |
-
)
|
54 |
)
|
55 |
|
56 |
# 4. Return the response text
|
57 |
return response.text
|
58 |
|
|
|
59 |
async def embedding_func(texts: list[str]) -> np.ndarray:
|
60 |
-
model = SentenceTransformer(
|
61 |
embeddings = model.encode(texts, convert_to_numpy=True)
|
62 |
return embeddings
|
63 |
|
|
|
64 |
rag = LightRAG(
|
65 |
working_dir=WORKING_DIR,
|
66 |
llm_model_func=llm_model_func,
|
@@ -72,11 +71,14 @@ rag = LightRAG(
|
|
72 |
)
|
73 |
|
74 |
file_path = "story.txt"
|
75 |
-
with open(file_path,
|
76 |
text = file.read()
|
77 |
|
78 |
rag.insert(text)
|
79 |
|
80 |
-
response = rag.query(
|
|
|
|
|
|
|
81 |
|
82 |
-
print
|
|
|
47 |
response = client.models.generate_content(
|
48 |
model="gemini-1.5-flash",
|
49 |
contents=[combined_prompt],
|
50 |
+
config=types.GenerateContentConfig(max_output_tokens=500, temperature=0.1),
|
|
|
|
|
|
|
51 |
)
|
52 |
|
53 |
# 4. Return the response text
|
54 |
return response.text
|
55 |
|
56 |
+
|
57 |
async def embedding_func(texts: list[str]) -> np.ndarray:
|
58 |
+
model = SentenceTransformer("all-MiniLM-L6-v2")
|
59 |
embeddings = model.encode(texts, convert_to_numpy=True)
|
60 |
return embeddings
|
61 |
|
62 |
+
|
63 |
rag = LightRAG(
|
64 |
working_dir=WORKING_DIR,
|
65 |
llm_model_func=llm_model_func,
|
|
|
71 |
)
|
72 |
|
73 |
file_path = "story.txt"
|
74 |
+
with open(file_path, "r") as file:
|
75 |
text = file.read()
|
76 |
|
77 |
rag.insert(text)
|
78 |
|
79 |
+
response = rag.query(
|
80 |
+
query="What is the main theme of the story?",
|
81 |
+
param=QueryParam(mode="hybrid", top_k=5, response_type="single line"),
|
82 |
+
)
|
83 |
|
84 |
+
print(response)
|