From 0fb0c2c61ef6f48089b941f0e185e51288fe01b0 Mon Sep 17 00:00:00 2001
From: shidong <shidong@jhsoft.cc>
Date: 星期二, 08 七月 2025 19:07:04 +0800
Subject: [PATCH] #2025/7/8 #增了读取知识库ID
---
qwen_thread.py | 164 ++++++++++++++++++++++++++++++++++++++++++++++--------
1 files changed, 140 insertions(+), 24 deletions(-)
diff --git a/qwen_thread.py b/qwen_thread.py
index 56ba944..85275f9 100644
--- a/qwen_thread.py
+++ b/qwen_thread.py
@@ -33,7 +33,7 @@
for i in range(max_workers):
model = AutoModelForVision2Seq.from_pretrained(
model_path,
- device_map="cuda:1",
+ device_map=f"cuda:{config.get('cuda')}",
trust_remote_code=True,
torch_dtype=torch.float16
).eval()
@@ -85,18 +85,33 @@
raise
def tark_do(self,res,ragurl,rag_mode,max_tokens):
- try :
+ try:
# 1. 浠庨泦鍚圓鑾峰彇鍚戦噺鍜屽厓鏁版嵁
is_waning = 0
- image_des = self.image_desc(f"{res['image_desc_path']}")
- self.logger.info(image_des)
+ is_desc = 2
+
+ # 鐢熸垚鍥剧墖鎻忚堪
+ image_des = self.image_desc(res['image_desc_path'])
+ risk_description = ""
+ suggestion = ""
+ # 鍥剧墖鎻忚堪鐢熸垚鎴愬姛
if image_des:
- # rule_text = self.get_rule(ragurl)
- is_waning = self.image_rule_chat(image_des,res['waning_value'],ragurl,rag_mode,max_tokens)
is_desc = 2
+ # 璋冪敤瑙勫垯鍖归厤鏂规硶,鍒ゆ柇鏄惁棰勮
+ is_waning = self.image_rule_chat(image_des, res['waning_value'], ragurl, rag_mode, max_tokens)
+ # 濡傛灉棰勮,鍒欑敓鎴愰殣鎮ｆ弿杩板拰澶勭悊寤鸿
+ if is_waning == 1:
+ # 鑾峰彇瑙勭珷鍒跺害鏁版嵁
+ filedata = self.get_filedata(res['waning_value'], ragurl)
+ # 鐢熸垚闅愭偅鎻忚堪
+ risk_description = self.image_rule_chat_with_detail(filedata, res['waning_value'], ragurl, rag_mode)
+ # 鐢熸垚澶勭悊寤鸿
+ suggestion = self.image_rule_chat_suggestion(res['waning_value'], ragurl, rag_mode)
+
else:
- is_waning = 0
is_desc = 3
+
+ # 鏁版嵁缁�
data = {
"id": res['id'],
"event_level_id": res['event_level_id'], # event_level_id
@@ -117,27 +132,31 @@
"image_path": res['image_path'], # image_path
"image_desc_path": res['image_desc_path'], # image_desc_path
"video_path": res['video_path'],
- "text_vector": res['text_vector']
+ "text_vector": res['text_vector'],
+ "risk_description": risk_description,
+ "suggestion": suggestion,
+ "knowledge_id": res['knowledge_id']
}
+
# 淇濆瓨鍒癿ilvus
image_id = self.collection.upsert(data).primary_keys
- if is_desc == 2:
- data = {
- "id": str(image_id[0]),
- "video_point_id": res['video_point_id'],
- "video_path": res["video_point_name"],
- "zh_desc_class": image_des,
- "detect_time": res['detect_time'],
- "image_path": f"{res['image_path']}",
- "task_name": res["task_name"],
- "event_level_name": res["event_level_name"],
- "rtsp_address": f"{res['video_path']}"
- }
- # 璋冪敤rag
- asyncio.run(self.insert_json_data(ragurl, data))
+ logging.info(image_id)
+ data = {
+ "id": str(image_id[0]),
+ "video_point_id": res['video_point_id'],
+ "video_path": res["video_point_name"],
+ "zh_desc_class": image_des,
+ "detect_time": res['detect_time'],
+ "image_path": f"{res['image_path']}",
+ "task_name": res["task_name"],
+ "event_level_name": res["event_level_name"],
+ "rtsp_address": f"{res['video_path']}"
+ }
+ # 璋冪敤rag
+ asyncio.run(self.insert_json_data(ragurl, data))
return image_id
except Exception as e:
- self.logger.info(f"绾跨▼锛氭墽琛屾ā鍨嬭В鏋愭椂鍑洪敊:浠诲姟锛歿e}")
+ self.logger.info(f"绾跨▼锛氭墽琛屾ā鍨嬭В鏋愭椂鍑洪敊:浠诲姟锛歿res['id']} :{e}")
return 0
def image_desc(self, image_path):
@@ -167,7 +186,7 @@
padding=True,
return_tensors="pt",
)
- inputs = inputs.to("cuda:1")
+ inputs = inputs.to(model.device)
current_time = datetime.now()
outputs = model.generate(**inputs,
max_new_tokens=300,
@@ -253,6 +272,103 @@
self.logger.info(f"绾跨▼锛氭墽琛岃鍒欏尮閰嶆椂鍑洪敊:{image_des, rule_text, ragurl, rag_mode,e}")
return None
+ # 闅愭偅鎻忚堪
+ def image_rule_chat_with_detail(self, filedata, rule_text, ollama_url, ollama_mode="qwen2.5vl:3b"):
+
+ # API璋冪敤
+ response = requests.post(
+ # ollama鍦板潃
+ url=f"{ollama_url}/chat",
+ json={
+ "prompt": "",
+ # 璇锋眰鍐呭
+ "messages": [
+ {
+ "role": "user",
+ "content": f"璇锋牴鎹绔犲埗搴{filedata}]\n鏌ユ壘[{rule_text}]鐨勫畨鍏ㄩ殣鎮ｆ弿杩帮紝涓嶈繘琛屾帹鐞嗗拰think銆傝繑鍥炰俊鎭皬浜�800瀛�"
+ }
+ ],
+ # 鎸囧畾妯″瀷
+ "llm_name": "qwen3:8b",
+ "stream": False, # 鍏抽棴娴佸紡杈撳嚭
+ "gen_conf": {
+ "temperature": 0.7, # 鎺у埗鐢熸垚闅忔満鎬�
+ "max_tokens": 800 # 鏈�澶ц緭鍑洪暱搴�
+ }
+ }
+ )
+ # 浠巎son鎻愬彇data瀛楁鍐呭
+ ret = response.json()["data"]
+ # result = response.json()
+ # ret = result.get("data") or result.get("message", {}).get("content", "")
+ # 绉婚櫎<think>鏍囩鍜屽唴瀹�
+ ret = re.sub(r'<think>.*?</think>', '', ret, flags=re.DOTALL)
+ # 瀛楃涓叉竻鐞�,绉婚櫎绌烘牸,鍒惰〃绗�,鎹㈣绗�,鏄熷彿
+ ret = ret.replace(" ", "").replace("\t", "").replace("\n", "").replace("**", "")
+ print(ret)
+ return ret
+
+ # 澶勭悊寤鸿
+ def image_rule_chat_suggestion(self, rule_text, ollama_url, ollama_mode="qwen2.5vl:3b"):
+ self.logger.info("----------------------------------------------------------------")
+ # 璇锋眰鍐呭
+ content = (
+ f"璇锋牴鎹繚瑙勫唴瀹筟{rule_text}]\n杩涜杩斿洖澶勭悊杩濊寤鸿锛屼笉杩涜鎺ㄧ悊鍜宼hink銆傝繑鍥炵簿鍑嗕俊鎭�")
+ response = requests.post(
+ # ollama鍦板潃
+ url=f"{ollama_url}/chat",
+ json={
+ # 鎸囧畾妯″瀷
+ "llm_name": "qwen3:8b",
+ "messages": [
+ {"role": "user", "content": content}
+ ],
+ "stream": False # 鍏抽棴娴佸紡杈撳嚭
+ }
+ )
+ # 浠巎son鎻愬彇data瀛楁鍐呭
+ ret = response.json()["data"]
+ # result = response.json()
+ # ret = result.get("data") or result.get("message", {}).get("content", "")
+ # 绉婚櫎<think>鏍囩鍜屽唴瀹�
+ ret = re.sub(r'<think>.*?</think>', '', ret, flags=re.DOTALL)
+ # 瀛楃涓叉竻鐞�,绉婚櫎绌烘牸,鍒惰〃绗�,鎹㈣绗�,鏄熷彿
+ ret = ret.replace(" ", "").replace("\t", "").replace("\n", "").replace("**", "")
+ print(ret)
+ return ret
+
+ # RAG鏈嶅姟鍙戦�佽姹�,鑾峰彇鐭ヨ瘑搴撳唴瀹�
+ def get_filedata(self, searchtext, ragurl):
+ search_data = {
+ # 鐭ヨ瘑搴撻泦鍚�
+ "collection_name": "smart_knowledge",
+ # 鏌ヨ鏂囨湰
+ "query_text": searchtext,
+ # 鎼滅储妯″紡
+ "search_mode": "hybrid",
+ # 鏈�澶氳繑鍥炵粨鏋�
+ "limit": 100,
+ # 璋冨瘑鍚戦噺鎼滅储鏉冮噸
+ "weight_dense": 0.7,
+ # 绋�鐤忓悜閲忔悳绱㈡潈閲�
+ "weight_sparse": 0.3,
+ # 绌哄瓧绗︿覆
+ "filter_expr": "",
+ # 鍙繑鍥� text 瀛楁
+ "output_fields": ["text"]
+ }
+ # 鍚� ragurl + "/search" 绔偣鍙戦�丳OST璇锋眰
+ response = requests.post(ragurl + "/search", json=search_data)
+ # 浠庡搷搴斾腑鑾峰彇'results'瀛楁
+ results = response.json().get('results')
+ # 鍒濆鍖� text
+ text = ""
+ # 閬嶅巻鎵�鏈夌粨鏋滆鍒�(rule)锛屽皢姣忔潯瑙勫垯鐨�'entity'涓殑'text'瀛楁鍙栧嚭.
+ for rule in results:
+ text = text + rule['entity'].get('text') + ";\n"
+
+ return text
+
async def insert_json_data(self, ragurl, data):
try:
data = {'collection_name': "smartrag", "data": data, "description": ""}
--
Gitblit v1.8.0