Gurjot Singh commited on
Commit
5b5f653
·
1 Parent(s): d5d0856

fixed linting errors

Browse files
Files changed (1) hide show
  1. examples/lightrag_gemini_demo.py +10 -8
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('all-MiniLM-L6-v2')
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, 'r') as file:
76
  text = file.read()
77
 
78
  rag.insert(text)
79
 
80
- response = rag.query(query="What is the main theme of the story?", param=QueryParam(mode="hybrid", top_k=5, response_type="single line"))
 
 
 
81
 
82
- print (response)
 
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)