Aniun commited on
Commit
4c2fab7
·
1 Parent(s): 625ef00
Files changed (7) hide show
  1. .gitattributes +1 -0
  2. .gitignore +173 -0
  3. app.py +346 -0
  4. database/init/init.faiss +3 -0
  5. database/init/init.pkl +3 -0
  6. deal_data.py +136 -0
  7. requirements.txt +9 -0
.gitattributes CHANGED
@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
33
  *.zip filter=lfs diff=lfs merge=lfs -text
34
  *.zst filter=lfs diff=lfs merge=lfs -text
35
  *tfevents* filter=lfs diff=lfs merge=lfs -text
 
 
33
  *.zip filter=lfs diff=lfs merge=lfs -text
34
  *.zst filter=lfs diff=lfs merge=lfs -text
35
  *tfevents* filter=lfs diff=lfs merge=lfs -text
36
+ *.faiss filter=lfs diff=lfs merge=lfs -text
.gitignore ADDED
@@ -0,0 +1,173 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # local
2
+ data/
3
+ data_all/
4
+ data_simple/
5
+ compose.ipynb
6
+ deal_data_org.py
7
+ search.ipynb
8
+
9
+ .env.local
10
+ .env
11
+
12
+ # Byte-compiled / optimized / DLL files
13
+ __pycache__/
14
+ *.py[cod]
15
+ *$py.class
16
+
17
+ # C extensions
18
+ *.so
19
+
20
+ # Distribution / packaging
21
+ .Python
22
+ build/
23
+ develop-eggs/
24
+ dist/
25
+ downloads/
26
+ eggs/
27
+ .eggs/
28
+ lib/
29
+ lib64/
30
+ parts/
31
+ sdist/
32
+ var/
33
+ wheels/
34
+ share/python-wheels/
35
+ *.egg-info/
36
+ .installed.cfg
37
+ *.egg
38
+ MANIFEST
39
+
40
+ # PyInstaller
41
+ # Usually these files are written by a python script from a template
42
+ # before PyInstaller builds the exe, so as to inject date/other infos into it.
43
+ *.manifest
44
+ *.spec
45
+
46
+ # Installer logs
47
+ pip-log.txt
48
+ pip-delete-this-directory.txt
49
+
50
+ # Unit test / coverage reports
51
+ htmlcov/
52
+ .tox/
53
+ .nox/
54
+ .coverage
55
+ .coverage.*
56
+ .cache
57
+ nosetests.xml
58
+ coverage.xml
59
+ *.cover
60
+ *.py,cover
61
+ .hypothesis/
62
+ .pytest_cache/
63
+ cover/
64
+
65
+ # Translations
66
+ *.mo
67
+ *.pot
68
+
69
+ # Django stuff:
70
+ *.log
71
+ local_settings.py
72
+ db.sqlite3
73
+ db.sqlite3-journal
74
+
75
+ # Flask stuff:
76
+ instance/
77
+ .webassets-cache
78
+
79
+ # Scrapy stuff:
80
+ .scrapy
81
+
82
+ # Sphinx documentation
83
+ docs/_build/
84
+
85
+ # PyBuilder
86
+ .pybuilder/
87
+ target/
88
+
89
+ # Jupyter Notebook
90
+ .ipynb_checkpoints
91
+
92
+ # IPython
93
+ profile_default/
94
+ ipython_config.py
95
+
96
+ # pyenv
97
+ # For a library or package, you might want to ignore these files since the code is
98
+ # intended to run in multiple environments; otherwise, check them in:
99
+ # .python-version
100
+
101
+ # pipenv
102
+ # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
103
+ # However, in case of collaboration, if having platform-specific dependencies or dependencies
104
+ # having no cross-platform support, pipenv may install dependencies that don't work, or not
105
+ # install all needed dependencies.
106
+ #Pipfile.lock
107
+
108
+ # poetry
109
+ # Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
110
+ # This is especially recommended for binary packages to ensure reproducibility, and is more
111
+ # commonly ignored for libraries.
112
+ # https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
113
+ #poetry.lock
114
+
115
+ # pdm
116
+ # Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
117
+ #pdm.lock
118
+ # pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
119
+ # in version control.
120
+ # https://pdm.fming.dev/latest/usage/project/#working-with-version-control
121
+ .pdm.toml
122
+ .pdm-python
123
+ .pdm-build/
124
+
125
+ # PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
126
+ __pypackages__/
127
+
128
+ # Celery stuff
129
+ celerybeat-schedule
130
+ celerybeat.pid
131
+
132
+ # SageMath parsed files
133
+ *.sage.py
134
+
135
+ # Environments
136
+ .env
137
+ .venv
138
+ env/
139
+ venv/
140
+ ENV/
141
+ env.bak/
142
+ venv.bak/
143
+
144
+ # Spyder project settings
145
+ .spyderproject
146
+ .spyproject
147
+
148
+ # Rope project settings
149
+ .ropeproject
150
+
151
+ # mkdocs documentation
152
+ /site
153
+
154
+ # mypy
155
+ .mypy_cache/
156
+ .dmypy.json
157
+ dmypy.json
158
+
159
+ # Pyre type checker
160
+ .pyre/
161
+
162
+ # pytype static type analyzer
163
+ .pytype/
164
+
165
+ # Cython debug symbols
166
+ cython_debug/
167
+
168
+ # PyCharm
169
+ # JetBrains specific template is maintained in a separate JetBrains.gitignore that can
170
+ # be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
171
+ # and can be added to the global gitignore or merged into this file. For a more nuclear
172
+ # option (not recommended) you can uncomment the following to ignore the entire idea folder.
173
+ #.idea/
app.py ADDED
@@ -0,0 +1,346 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import time
3
+ import json
4
+ import asyncio
5
+ import gradio as gr
6
+
7
+ # set the env
8
+ from dotenv import load_dotenv
9
+ load_dotenv()
10
+
11
+ # get the root path of the project
12
+ current_file_path = os.path.dirname(os.path.abspath(__file__))
13
+ root_path = os.path.abspath(current_file_path)
14
+
15
+ from textwrap import dedent
16
+ from langchain_openai import ChatOpenAI
17
+ from langchain_openai import OpenAIEmbeddings
18
+ from langchain_community.vectorstores import FAISS
19
+ from langchain_core.prompts import ChatPromptTemplate
20
+
21
+ class OurLLM:
22
+ def __init__(self, model="gpt-4o"):
23
+ '''
24
+ params:
25
+ model: str,
26
+ 模型名称 ["GLM-4-Flash", "GLM-4V-Flash",
27
+ "gpt-4o-mini", "gpt-4o", "o1-mini",
28
+ "gemini-1.5-flash-002", "gemini-1.5-pro-002",
29
+ "Qwen/Qwen2.5-7B-Instruct", "Qwen/Qwen2.5-Coder-7B-Instruct"]
30
+ '''
31
+
32
+ self.model_name = model
33
+
34
+ OPENAI_API_KEY = os.getenv('OPENAI_API_KEY')
35
+ OPENAI_API_KEY_DF = os.getenv('OPENAI_API_KEY_DF', OPENAI_API_KEY)
36
+ OPENAI_API_KEY_AZ = os.getenv('OPENAI_API_KEY_AZ', OPENAI_API_KEY)
37
+ OPENAI_API_KEY_CD = os.getenv('OPENAI_API_KEY_CD')
38
+ OPENAI_API_KEY_O1 = os.getenv('OPENAI_API_KEY_O1')
39
+ OPENAI_API_KEY_GLM = os.getenv('OPENAI_API_KEY_GLM')
40
+ OPENAI_API_KEY_SC = os.getenv('OPENAI_API_KEY_SC')
41
+
42
+ OPENAI_BASE_URL = os.getenv('OPENAI_BASE_URL')
43
+ OPENAI_BASE_URL_GLM = os.getenv('OPENAI_BASE_URL_GLM')
44
+ OPENAI_BASE_URL_SC = os.getenv('OPENAI_BASE_URL_SC')
45
+
46
+ # 创建 API Key 映射
47
+ apiKeyMap = {
48
+ 'gemini': {"base_url": OPENAI_BASE_URL, "api_key": OPENAI_API_KEY_DF},
49
+ 'gpt': {"base_url": OPENAI_BASE_URL, "api_key": OPENAI_API_KEY_AZ},
50
+ 'o1': {"base_url": OPENAI_BASE_URL, "api_key": OPENAI_API_KEY_O1},
51
+ 'claude': {"base_url": OPENAI_BASE_URL, "api_key": OPENAI_API_KEY_CD},
52
+ 'glm': {"base_url": OPENAI_BASE_URL_GLM, "api_key": OPENAI_API_KEY_GLM},
53
+ 'qwen': {"base_url": OPENAI_BASE_URL_SC, "api_key": OPENAI_API_KEY_SC},
54
+ }
55
+
56
+ for name, info in apiKeyMap.items():
57
+ if name in model.lower():
58
+ self.base_url = info["base_url"]
59
+ self.api_key = info["api_key"]
60
+ break
61
+ assert self.base_url is not None, f"Base URL not found for model: {model}"
62
+ assert self.api_key is not None, f"API key not found for model: {model}"
63
+
64
+ chat_prompt = ChatPromptTemplate.from_messages(
65
+ [
66
+ ("system", "{system_prompt}"),
67
+ ("human", "{input}"),
68
+ # ("ai", "{chat_history}"),
69
+ ]
70
+ )
71
+ self.chat_prompt = chat_prompt
72
+ self.llm = self.get_llm(model)
73
+
74
+ def clean_json(self, s):
75
+ return s.replace("```json", "").replace("```", "").strip()
76
+
77
+ def get_system_prompt(self, mode="assistant"):
78
+ prompt_map = {
79
+ "assistant": dedent("""
80
+ 你是一个智能助手,擅长用简洁的中文回答用户的问题。
81
+ 请确保你的回答准确、清晰、有条理,并且符合中文的语言习惯。
82
+ 重要提示:
83
+ 1. 回答要简洁明了,避免冗长
84
+ 2. 使用适当的专业术语
85
+ 3. 保持客观中立的语气
86
+ 4. 如果不确定,要明确指出
87
+ """),
88
+ # search
89
+ "keyword_expand": dedent("""
90
+ 你是一个搜索关键词扩展专家,擅长将用户的搜索意图转化为多个相关的搜索词或短语。
91
+ 用户会输入一段描述他们搜索需求的文本,请你生成与之相关的关键词列表。
92
+ 你需要返回一个可以直接被 json 库解析的响应,包含以下内容:
93
+ {
94
+ "keywords": [关键词列表],
95
+ }
96
+ 重要提示:
97
+ 1. 关键词应该包含同义词、近义词、上位词、下位词
98
+ 2. 短语要体现不同的表达方式和组合
99
+ 3. 描述句子要涵盖不同的应用场景和用途
100
+ 4. 所有内容必须与原始搜索意图高度相关
101
+ 5. 扩展搜索意图到相关的应用场景和工具,例如:
102
+ - 如果搜索"PDF转MD",应包含PDF内容提取、PDF解析工具、PDF数据处理等
103
+ - 如果搜索"图片压缩",应包含批量压缩工具、图片格式转换等
104
+ - 如果搜索"代码格式化",应包含代码美化工具、语法检查器、代码风格统一等
105
+ - 如果搜索"文本翻译",应包含机器翻译API、多语言翻译工具、离线翻译软件等
106
+ - 如果搜索"数据可视化",应包含图表生成工具、数据分析库、交互式图表等
107
+ - 如果搜索"网络爬虫",应包含数据采集框架、反爬虫绕过、数据解析工具等
108
+ - 如果搜索"API测试",应包含接口测试工具、性能监控、自动化测试框架等
109
+ 6. 所有内容主要使用英文表达,并对部分关键词添加额外的中文表示
110
+ 7. 返回内容不要使用任何 markdown 格式 以及任何特殊字符
111
+ """),
112
+ "zh2en": dedent("""
113
+ 你是一个专业的中译英翻译专家,尤其擅长学术论文的翻译工作。
114
+ 请将用户提供的中文内容翻译成地道、专业的英文。
115
+
116
+ 重要提示:
117
+ 1. 使用学术论文常用的表达方式和术语
118
+ 2. 保持专业、正式的语气
119
+ 3. 确保译文的准确性和流畅性
120
+ 4. 对专业术语进行准确翻译
121
+ 5. 遵循英文学术写作的语法规范
122
+ 6. 保持原文的逻辑结构
123
+ 7. 适当使用学术论文常见的过渡词和连接词
124
+ 8. 如遇到模糊的表达,选择最符合学术上下文的翻译
125
+ 9. 避免使用口语化或非正式的表达
126
+ 10. 注意时态和语态的准确使用
127
+ """),
128
+ "github_score": dedent("""
129
+ 你是一个语义匹配评分专家,擅长根据用户需求和仓库描述进行语义匹配度评分。
130
+ 用户会输入两部分内容:
131
+ 1. 用户的具体需求描述
132
+ 2. 多个仓库的描述列表(以1,2,3等数字开头)
133
+
134
+ 请你仔细分析用户需求,并对每个仓库进行评分。
135
+ 确保返回一个可以直接被 json 库解析的响应,包含以下内容:
136
+ {
137
+ "indices": [仓库编号列表,按分数从高到低],
138
+ "scores": [编号对应的匹配度评分列表,0-100的整数,表示匹配程度]
139
+ }
140
+
141
+ 重要提示:
142
+ 1. 评分范围为0-100的整数,高于60分表示具有明显相关性
143
+ 2. 评分要客观反映仓库与需求的契合度
144
+ 3. 只返回评分大于 60 的仓库
145
+ 4. 返回内容不要使用任何 markdown 格式 以及任何特殊字符
146
+ """)
147
+ }
148
+ return prompt_map[mode]
149
+
150
+ def get_llm(self, model="gpt-4o-mini"):
151
+ '''
152
+ params:
153
+ model: str, 模型名称 ["gpt-4o-mini", "gpt-4o", "o1-mini", "gemini-1.5-flash-002"]
154
+ '''
155
+ llm = ChatOpenAI(
156
+ model=model,
157
+ base_url=self.base_url,
158
+ api_key=self.api_key,
159
+ )
160
+ print(f"Init model {model} successfully!")
161
+ return llm
162
+
163
+ def ask_question(self, question, system_prompt=None):
164
+ # 1. 获取系统提示
165
+ if system_prompt is None:
166
+ system_prompt = self.get_system_prompt()
167
+
168
+ # 2. 生成聊天提示
169
+ prompt = self.chat_prompt.format(input=question, system_prompt=system_prompt)
170
+ config = {
171
+ "configurable": {"response_format": {"type": "json_object"}}
172
+ }
173
+
174
+ # 3. 调用 LLM 进行回答
175
+ for _ in range(10):
176
+ try:
177
+ response = self.llm.invoke(prompt, config=config)
178
+ response.content = self.clean_json(response.content)
179
+ return response
180
+ except Exception as e:
181
+ print(e)
182
+ time.sleep(10)
183
+ continue
184
+ print(f"Failed to call llm for prompt: {prompt[0:10]}")
185
+ return None
186
+
187
+ async def ask_questions_parallel(self, questions, system_prompt=None):
188
+ # 1. 获取系统提示
189
+ if system_prompt is None:
190
+ system_prompt = self.get_system_prompt()
191
+
192
+ # 2. 定义异步函数
193
+ async def call_llm(prompt):
194
+ for _ in range(10):
195
+ try:
196
+ response = await self.llm.ainvoke(prompt)
197
+ response.content = self.clean_json(response.content)
198
+ return response
199
+ except Exception as e:
200
+ print(e)
201
+ await asyncio.sleep(10)
202
+ continue
203
+ print(f"Failed to call llm for prompt: {prompt[0:10]}")
204
+ return None
205
+
206
+ # 3. 构建 prompt
207
+ prompts = [self.chat_prompt.format(input=question, system_prompt=system_prompt) for question in questions]
208
+
209
+ # 4. 异步调用
210
+ tasks = [call_llm(prompt) for prompt in prompts]
211
+ results = await asyncio.gather(*tasks)
212
+
213
+ return results
214
+
215
+ class RepoSearch:
216
+ def __init__(self):
217
+ db_path = os.path.join(root_path, "database", "init")
218
+ embeddings = OpenAIEmbeddings(api_key=os.getenv("OPENAI_API_KEY"),
219
+ base_url=os.getenv("OPENAI_BASE_URL"),
220
+ model="text-embedding-3-small")
221
+
222
+ assert os.path.exists(db_path), f"Database not found: {db_path}"
223
+ self.vector_db = FAISS.load_local(db_path, embeddings,
224
+ index_name="init",
225
+ allow_dangerous_deserialization=True)
226
+
227
+ def search(self, query, k=10):
228
+ '''
229
+ name + description + html_url + topics
230
+ '''
231
+ results = self.vector_db.similarity_search(query + " technology", k=k)
232
+
233
+ simple_str = ""
234
+ simple_list = []
235
+ for i, doc in enumerate(results):
236
+ content = json.loads(doc.page_content)
237
+ metadata = doc.metadata
238
+ if content["description"] is None:
239
+ content["description"] = ""
240
+ # desc = content["description"] if len(content["description"]) < 300 else content["description"][:300] + "..."
241
+ simple_str += f"\t**{i+1}. {content['name']}** || {content['description']}\n" # 用于大模型匹配
242
+ simple_list.append({
243
+ "name": content["name"],
244
+ "description": content["description"],
245
+ **metadata, # 解包所有 metadata 字段
246
+ })
247
+
248
+ return simple_str, simple_list
249
+
250
+ def main():
251
+ search = RepoSearch()
252
+ llm = OurLLM(model="gpt-4o")
253
+
254
+ def respond(
255
+ prompt: str,
256
+ history,
257
+ is_llm_filter: bool = False,
258
+ is_keyword_expand: bool = False,
259
+ match_num: int = 40
260
+ ):
261
+ # 1. 初始化历史记录
262
+ if not history:
263
+ history = [{"role": "system", "content": "You are a friendly chatbot"}]
264
+ history.append({"role": "user", "content": prompt})
265
+ response = {"role": "assistant", "content": ""}
266
+ yield history
267
+
268
+ # 2. 扩展用户问题关键词
269
+ if is_keyword_expand:
270
+ response["content"] = "开始扩展关键词..."
271
+ yield history + [response]
272
+
273
+ query = llm.ask_question(prompt, system_prompt=llm.get_system_prompt("keyword_expand")).content
274
+ prompt = ", ".join(json.loads(query)["keywords"])
275
+
276
+ # 3. 语义向量匹配
277
+ response["content"] = "开始语义向量匹配..."
278
+ yield history + [response]
279
+ match_str, simple_list = search.search(prompt, match_num)
280
+
281
+ # 4. 通过 LLM 评分得到最匹配的仓库索引
282
+ if not is_llm_filter:
283
+ simple_strs = [f"\t**{i+1}. {repo['name']}** [✨ {repo['star_count'] // 1000}k] || **Description:** {repo['description']} || **Url:** {repo['html_url']} \n" for i, repo in enumerate(simple_list)]
284
+ response["content"] = "".join(simple_strs)
285
+ yield history + [response]
286
+ else:
287
+ response["content"] = "开始通过 LLM 评分得到最匹配的仓库..."
288
+ yield history + [response]
289
+
290
+ query = ' ## 用户需要的仓库内容:' + prompt + '\n ## 搜索结果列表:' + match_str
291
+ out = llm.ask_question(query, system_prompt=llm.get_system_prompt("github_score")).content
292
+ matched_index = json.loads(out)["indices"]
293
+
294
+ # 5. 通过索引得到最匹配的仓库
295
+ result = [simple_list[idx-1] for idx in matched_index]
296
+ simple_strs = [f"\t**{i+1}. {repo['name']}** [✨ {repo['star_count'] // 1000}k] || **Description:** {repo['description']} || **Url:** {repo['html_url']} \n" for i, repo in enumerate(result)]
297
+ response["content"] = "".join(simple_strs)
298
+ yield history + [response]
299
+
300
+ with gr.Blocks() as demo:
301
+ gr.Markdown("## Github semantic search (基于语义的 github 仓库搜索) 🌐")
302
+
303
+ with gr.Row():
304
+ with gr.Column(scale=1):
305
+ # 添加控制参数
306
+ llm_filter = gr.Checkbox(
307
+ label="使用LLM过滤结果",
308
+ value=False,
309
+ info="是否使用 LLM 对搜索结果进行二次过滤"
310
+ )
311
+ keyword_expand = gr.Checkbox(
312
+ label="扩展关键词搜索",
313
+ value=False,
314
+ info="是否使用 LLM 扩展搜索关键词"
315
+ )
316
+ match_number = gr.Slider(
317
+ minimum=10,
318
+ maximum=100,
319
+ value=40,
320
+ step=10,
321
+ label="语义匹配数量",
322
+ info="进行语义匹配后返回的仓库数量,若使用 LLM 过滤,建议适当增加数量"
323
+ )
324
+
325
+ with gr.Column(scale=3):
326
+ chatbot = gr.Chatbot(
327
+ label="Agent",
328
+ type="messages",
329
+ avatar_images=(None, "https://img1.baidu.com/it/u=2193901176,1740242983&fm=253&fmt=auto&app=138&f=JPEG?w=500&h=500"),
330
+ height="65vh"
331
+ )
332
+ prompt = gr.Textbox(max_lines=2, label="Chat Message")
333
+
334
+ # 更新submit调用,包含新的参数
335
+ prompt.submit(
336
+ respond,
337
+ [prompt, chatbot, llm_filter, keyword_expand, match_number],
338
+ [chatbot]
339
+ )
340
+ prompt.submit(lambda: "", None, [prompt])
341
+
342
+ demo.launch(share=False)
343
+
344
+
345
+ if __name__ == "__main__":
346
+ main()
database/init/init.faiss ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:007ee89904bde801789ec8fd0c3b56dc1d6b2d684d3c4c80f808cf1614c38ad7
3
+ size 279736365
database/init/init.pkl ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:c3e3d8fd7c484feab60214829cf3e9f1edbb99e48d6453aed86a6899223bf379
3
+ size 19696870
deal_data.py ADDED
@@ -0,0 +1,136 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import json
3
+ import asyncio
4
+ import requests
5
+
6
+ from tqdm import tqdm
7
+ from dotenv import load_dotenv
8
+ load_dotenv()
9
+
10
+ from langchain_openai import ChatOpenAI
11
+ from langchain_core.prompts import ChatPromptTemplate
12
+ from langchain_core.documents import Document
13
+ from langchain_openai import OpenAIEmbeddings
14
+ from langchain_community.vectorstores import FAISS
15
+
16
+ # 获取当前目录根路径
17
+ current_file_path = os.path.dirname(os.path.abspath(__file__))
18
+ root_path = os.path.abspath(current_file_path)
19
+ data_path = os.path.join(root_path, "data_simple")
20
+ db_path = os.path.join(root_path, "database", "init")
21
+
22
+ # 1. 根据 star 数量区间获取 GitHub 仓库,同时根据 star 数量从多到少排序(闭区间)并保存 GitHub 仓库
23
+ def get_top_repo_by_star(per_page=1000, page=1, min_star_num=0, max_star_num=500000):
24
+ query = f'stars:{min_star_num}..{max_star_num} pushed:>2021-01-01'
25
+ sort = 'stars'
26
+ order = 'desc'
27
+ search_url = f'{os.getenv('GITHUB_API_URL')}/search/repositories?q={query}&sort={sort}&order={order}&per_page={per_page}&page={page}'
28
+ headers = {"Authorization": f"token {os.getenv('GITHUB_TOKEN')}"}
29
+
30
+ response = requests.get(search_url, headers=headers)
31
+ if response.status_code == 200:
32
+ total_count = response.json()['total_count']
33
+ total_page = total_count // per_page + 1
34
+ print(f"Total page: {total_page}, current page: {page}")
35
+ if response.json()['incomplete_results']: print("Incomplete results")
36
+ return response.json()['items'], response.json()['items'][-1]['stargazers_count'], total_count
37
+ else:
38
+ print(f"Failed to retrieve repositories: {response.status_code}")
39
+ print("")
40
+ # 直接退出
41
+ exit(1)
42
+
43
+ def save_repo_by_star(max_star=500000):
44
+ # github 限制每次请求最多得到 100 个仓库,因此 page 固定为 1
45
+ top_repositories, max_star, count = get_top_repo_by_star(per_page=1000, page=1, min_star_num=1000, max_star_num=max_star)
46
+
47
+ for i, repo in enumerate(top_repositories):
48
+ owner = repo['owner']['login']
49
+ name = repo['name']
50
+ unique_id = f"{name} -- {owner}"
51
+ stars = repo['stargazers_count']
52
+ print(f"Repository {i}: {name}, Stars: {stars}")
53
+
54
+ # 存储为 json 格式
55
+ with open(os.path.join(data_path, f'{unique_id}.json'), 'w') as f:
56
+ json.dump(repo, f, indent=4)
57
+
58
+ if count < 100: exit(1)
59
+
60
+ return max_star
61
+
62
+ def main_repo():
63
+ max_star = 500000 # 最多 star 的仓库有 500k
64
+ num = 1
65
+ while True:
66
+ print("=" * 50)
67
+ print(f"Round {num}, Max star: {max_star}")
68
+ max_star = save_repo_by_star(max_star)
69
+ num += 1
70
+
71
+ # 2. 将数据转换为向量
72
+ async def create_vector_db(docs, embeddings, batch_size=800):
73
+ # 初始化第一批数据
74
+ vector_db = await FAISS.afrom_documents(docs[0:batch_size], embeddings)
75
+ if len(docs) < batch_size: return vector_db
76
+
77
+ # 创建任务x``
78
+ tasks = []
79
+ for start_idx in range(batch_size, len(docs), batch_size):
80
+ end_idx = min(start_idx + batch_size, len(docs))
81
+ tasks.append(FAISS.afrom_documents(docs[start_idx:end_idx], embeddings))
82
+
83
+ # 执行任务
84
+ results = await asyncio.gather(*tasks)
85
+
86
+ # 合并结果
87
+ for temp_db in results:
88
+ vector_db.merge_from(temp_db)
89
+ return vector_db
90
+
91
+ async def main_convert_to_vector():
92
+ # 读取文件
93
+ files = os.listdir(data_path)
94
+
95
+ # 构建 document
96
+ docs = []
97
+ for file in tqdm(files):
98
+ if not file.endswith(".json"): continue
99
+ with open(os.path.join(data_path, file), "r", encoding="utf-8") as f:
100
+ data = json.load(f)
101
+
102
+ content_map = {
103
+ "name": data["name"],
104
+ "description": data["description"],
105
+ }
106
+ content = json.dumps(content_map)
107
+ doc = Document(page_content=content, metadata={"html_url": data["html_url"],
108
+ "topics": data["topics"],
109
+ "created_at": data["created_at"],
110
+ "updated_at": data["updated_at"],
111
+ "star_count": data["stargazers_count"]})
112
+ docs.append(doc)
113
+ print(f"Total {len(docs)} documents.")
114
+
115
+ # 初始化 Embedding 实例
116
+ embeddings = OpenAIEmbeddings(api_key=os.getenv("OPENAI_API_KEY"),
117
+ base_url=os.getenv("OPENAI_BASE_URL"),
118
+ model="text-embedding-3-small")
119
+ print("Embedding model success: text-embedding-3-small")
120
+
121
+ # 文档嵌入
122
+ if os.path.exists(os.path.join(db_path, "init.faiss")):
123
+ vector_db = FAISS.load_local(db_path, embeddings=embeddings,
124
+ index_name="init",
125
+ allow_dangerous_deserialization=True)
126
+ else:
127
+ vector_db = await create_vector_db(docs, embeddings=embeddings)
128
+ vector_db.save_local(db_path, index_name="init")
129
+ return vector_db
130
+
131
+ if __name__ == "__main__":
132
+ # 1. 获取仓库信息
133
+ # main_repo()
134
+
135
+ # 2. 构建向量数据库
136
+ asyncio.run(main_convert_to_vector())
requirements.txt ADDED
@@ -0,0 +1,9 @@
 
 
 
 
 
 
 
 
 
 
1
+ langchain_community
2
+ langchain_core
3
+ langchain_openai
4
+
5
+ faiss-cpu
6
+ tqdm
7
+ python-dotenv
8
+
9
+ gradio