from transformers import AutoModel, AutoTokenizer
from logger_config import set_logger
import torch

logger = set_logger('main_log', 'logs/default_log.log')

# 检查GPU是否可用
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 加载模型和分词器
# model_name = "intfloat/e5-large-v2"
# model_name = "BAAI/bge-reranker-base"
# model_name = "BAAI/bge-reranker-v2-m3"
model_name = "BAAI/bge-reranker-large"

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name).to(device)

def embed_text(text):
    inputs = tokenizer(text, return_tensors='pt', truncation=True, padding=True).to(device)
    with torch.no_grad():
        embeddings = model(**inputs).last_hidden_state.mean(dim=1)
    return embeddings

def calculate_similarity(text1, text2):
    embeddings1 = embed_text(text1)
    embeddings2 = embed_text(text2) 
    similarity = torch.nn.functional.cosine_similarity(embeddings1, embeddings2)
    return similarity.item()

# 示例文本
txt_1 = "软件平台的功能有哪些?"

# txt_2 = "软件平台功能软件平台，采用统一门户登录，统一角色、用户管理，不同用户分级、分权限使用。包含以下主要功能模块：1、实时视频、定位监控。可实时显示船舶的最新抓拍画面（抓拍间隔不低于30秒），点击图片能瞬时定位该船舶所在位置和基本信息，并可查看船舶轨迹和视频。2、视频展示。可对船舶进行分类展示，全部、在线、航行3种类别，并在主页上显示船舶总数、实时在线船舶总数、实时航行船舶总数；可打开实时视频监控，并在地图上定位船舶所在位置。3、驾驶员人脸识别。通过自动比对驾驶员信息库，识别驾驶员信息。4、甲板逗留人员识别。识别客船船外甲板上是否有乘客逗留。5、乘客人数统计。统计船舱内人数。6、语音告警。当客船发生超载、无证驾驶、乘客在甲板上逗留等情况，监管部门可通过平台语音广播提醒驾驶员。7、异常报警。当终端出现视频丢失、定位异常、本地存储异常的现象时，向平台发送报警信息；进入到异常报警统计报表后，其结果可进行查询和导出操作。8.船舶里程统计。船舶日行驶里程统计，将单船的某月份的行驶里程数按日以图表形式展现；船舶里程统计报表，自定义时段统计单船里程信息，其结果可进行查询和导出操作；船舶总里程统计报表，按辖区统计所有船舶总里程统，其结果可进行查询和导出操作。9、视频回放。可回放船舶的本地存储视频。10、轨迹回放。可对船舶轨迹进行回放。11、船舶历史监控图像。历史监控图像显示船舶的定时抓拍数据(抓拍间隔不低于30秒），图片可按船舶名称、时间进行查询。12、运维统计。按月、年统计省、市、县三级监控点的正常情况。手机APP1、实时视频功能。用户根据权限浏览实时视频，并可回放、抓图、录像、截图等。2、收藏夹功能。可对重点关注的船舶、码头视频进行单独分组收藏。3、地图展示功能。输入船舶名称，可以在地图上定位船舶所在位置。4、消息展示功能。支持实时展示触发的事件信息，并支持对历史事件安报警时间、类型等进行查询数据存储要求1、船舶视频前端存储不低于15天；码头视频前端存储不低于7天，后台存储不低于30天。支持H.264、H.265标准，采用H.265标准。视频图像分辨率不低于1920×1080，帧率不低于25fps；录像支持信息叠加。2、截图等图片数据后台存储不少于30天。3、轨迹数据、预警数据等存储不低于3年。软件平台性能1、系统命令执行时间。系统命令执行时间≤1s。2、实时信号延时性。系统为200个用户提供并发数据服务时，实时信号延时≤2s。3、查询响应时间。系统为200个用户提供并发数据服务时，历史数据统计查询时平均响应时间≤5s。4、支持用户数。支持至少500个用户同时在线使用。系统安全可视化监管平台数据库、通讯程序和WEB页面应分别部署在不同的服务器上，并安装必要的安全设备，加强数据安全管理。服务器、存储、网络安全设施等，采用云主机、云存储、云安全卫士、云专线。系统对接提供系统视频、定位、预警等数据的共享和对接。如：视频数据接入湖南省综合交通运输信息平台；视频数据、定位数据、航次客流量数据、驾驶员打卡数据等接入智慧水运综合监管平台，综合监管平台为客渡通小程序提供数据。中电信数智科技有限公司 1第三部分 技术文件"
# txt_2 = "平台软件主要功能实时视频、定位监控可实时显示船舶的最新抓拍画面（抓拍间隔不低于30秒）。实时监控画面点击图片能瞬时定位该船舶所在位置和基本信息，并可查看船舶轨迹和视频。定位监控画面视频展示可对船舶进行分类展示，全部、在线、航行3种类别，并在主页上显示船舶总数、实时在线船舶总数、实时航行船舶总数；可打开实时视频监控，并在地图上定位船舶所在位置。视频分类展示地图定位驾驶员人脸识别系统支持驾驶员身份识别，并且提前将驾驶员照片录入识别库（识别库支持500张图片），当系统检测到驾驶员不是识别库中已经录入的人员或者驾驶员发生更换时，系统作出判断提醒，发出报警传输到后台平台，平台做出相关预警以及记录。甲板逗留人员识别识别客船船外甲板上是否有乘客逗留。船载甲板通过安装视频相机覆盖。基于前端摄像机的智能识别算法，判断甲板是否存在行驶过程中人员逗留的情况，发出报警传输到后台平台，平台做出相关预警以及记录。乘客人数统计船载船舱通过多个客流相机完整覆盖。基于船载主机通过固定时间抓拍船舱内乘客图片，图片传输到后台平台，平台分析出船载具体人数详情。具体功能如下：支持图片抓拍的客流量统计；支持多通道客流数据汇总；支持单通道客流阈值设置，并可联动前端视频进行弹窗预警、语音报警、电子邮件、短信通知；支持区域内人数阈值设置及弹窗预警、语音报警；支持按照时间段、监控点等进行报警信息查询，并支持报警信息导出，便于事后查询；支持客流报表的生成（柱状图，折线图，时/天/周/月/年报表）和导出(柱状图、折线图以图片格式导出、报表以EXCEL形式导出）。视频智能识别信息显示及统计针对驾驶员人脸识别、甲板逗留人员识别及乘客人数统计等视频设备发现异常情况的报警显示及统计功能。通过对接设备接口，实现对前端设备驾驶员人脸识别情况、甲板逗留人员识别情况及乘客人数统计情况的数据采集与存储功能；实现上述三类数据预警阀值的设置功能，并根据设置的预警阀值，实现针对前端摄像头分析数据的实时声光预警功能（如通过设置船舶最大载客量，从而实现超载预警功能以及针对预警信息的处置台账记录功能）；根据系统产生的预警数据、处置记录台账等实现相关预警统计、处置情况分析以及统计数据导出功能。语音告警当客船发生超载、无证驾驶、乘客在甲板上逗留等情况，监管部门可通过平台语音广播提醒驾驶员。提醒方式包含TTS文字转语音及直接喊话。通过设备接口实现。异常报警"
txt_2 = "目录1.项目概述61.1.项目背景及建设目的61.2.系统建设需求61.2.1.船舶移动视频监控点技术要求61.2.2.码头固定监控点技术要求71.2.3.太阳能供电技术要求71.2.4.监控中心管理终端PC71.2.5.软件平台功能81.2.6.手机APP81.2.7.数据存储要求91.2.8.软件平台性能91.2.9.系统安全91.2.10.系统对接92.总体设计102.1.设计原则102.2.指导思想102.3.设计标准112.4.设计依据142.5.关键技术172.5.1.基于SOA的设计平台框架172.5.2.基于J2EE的技术开发路线182.5.3.基于OGC规范的电子江图Web服务192.5.4.基于HTML5的WEBSOCKET通讯技术202.6.系统架构设计212.7.软件配置环境232.7.1.操作系统232.7.2.数据库平台232.8.非功能性设计252.8.1.系统运行性能252.8.2.系统安全性252.8.3.系统可靠性252.8.4.系统易用性253.前端系统设计273.1.系统组成273.2.船载主机273.3.客船摄像机283.4.太阳能供电303.5.码头固定摄像机313.5.1.设计说明313.5.2.监控点设备组成313.5.3.监控中心管理终端PC324.网络传输系统设计334.1.无线视频传输简介334.2.4G无线网络传输优势334.3.网络系统组网设计345.公有云平台（天翼云）355.1.总体架构355.2.配置说明355.3.软件架构365.4.产品功能375.4.1.云平台385.4.2.云存储405.4.3.云点播（基于对象存储的视频存储播放综合解决方案）435.4.4.云安全卫士445.5.重难点分析465.5.1.数据对接解决方案465.5.2.数据安全解决方案486.PON云专网方案716.1.方案背景716.2.客户业务需求726.2.1.网络线路保证需求726.2.2.网络安全需求726.3.方案简述726.3.1.接入拓扑736.3.2.接入带宽（速率）736.3.3.对客户接入的相关要求746.3.4.QoS等级服务746.3.5.安全方案747.后端平台设计757.1.系统功能结构757.2.平台软件主要功能757.2.1.实时视频、定位监控757.2.2.视频展示767.2.3.驾驶员人脸识别77"

similarity_score = calculate_similarity(txt_1, txt_2)
logger.info(f";text_1 : {txt_1} ; text_2 : {txt_2[:15]}... ; Similarity score: {similarity_score} ———[{model_name}] ")

print(f"Similarity score: {similarity_score}")