Update moderations.py

moderations.py

@@ -1,7 +1,4 @@
 from gevent import pywsgi
-import dotenv
-dotenv.load_dotenv(override=True)
-
 import sys
 import time
 import argparse
@@ -10,9 +7,14 @@ from typing import Union
 from pydantic import BaseModel
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import torch
+import os
 import openedai
 import numpy as np
 import asyncio
+from urllib.parse import urlparse
+import nacos
+import configparser
+
 
 
 app = openedai.OpenAIStub()
@@ -35,7 +37,7 @@ labels = ['hate',
 
 label2id = {l:i for i, l in enumerate(labels)}
 id2label = {i:l for i, l in enumerate(labels)}
-model_name = "/root/autodl-tmp/duanyu027/moderation_0628"
+model_name = "duanyu027/moderation_0628"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=len(labels),id2label=id2label, label2id=label2id, problem_type = "multi_label_classification")
 model.to(device)
@@ -44,6 +46,38 @@ model.eval()
 #	model, {torch.nn.Linear}, dtype=torch.qint8
 #)
 torch.set_num_threads(1)
+def register_service(client,service_name,service_ip,service_port,cluster_name,health_check_interval,weight,http_proxy,domain,protocol,direct_domain):
+    try:
+        # 初始化 metadata
+        metadata = {}
+        
+        # 如果 http_proxy 为 True，添加额外的 metadata 键值对
+        if http_proxy:
+            metadata["http_proxy"] = True
+            if direct_domain:
+                metadata["domain"] = f"{protocol}://{service_ip}:{service_port}"
+            else:
+                metadata["domain"] = f"{domain}/port/{service_port}"
+        else:
+            metadata["http_proxy"] = False
+            metadata["domain"] = f"{protocol}://{service_ip}:{service_port}"
+        response = client.add_naming_instance(
+            service_name,
+            service_ip,
+            service_port,
+            cluster_name,
+            weight,
+            metadata,
+            enable=True,
+            healthy=True,
+            ephemeral=True,
+            heartbeat_interval=health_check_interval
+        )
+        return response
+    except Exception as e:
+        print(f"Error registering service to Nacos: {e}")
+        return True
+        
 class ModerationsRequest(BaseModel):
 	model: str = "text-moderation-latest" # or "text-moderation-stable"
 	input: Union[str, list[str]]
@@ -121,11 +155,41 @@ async def predict(text, model, tokenizer):
 # Main
 if __name__ == "__main__":
 
-	args = parse_args(sys.argv[1:])
-	# start API
-	print(f'Starting moderations[{device}] API on {args.host}:{args.port}', file=sys.stderr)
-	app.register_model('text-moderations-latest', 'text-moderations-stable')
-	app.register_model('text-moderations-005', 'text-moderations-ifmain')
+    # 创建配置解析器
+    config = configparser.ConfigParser()
+    # 读取配置文件
+    if not config.read('config.ini'):
+        raise RuntimeError("配置文件不存在")
+    # Nacos server and other configurations
+    NACOS_SERVER = config['nacos']['nacos_server']
+    NAMESPACE = config['nacos']['namespace']
+    CLUSTER_NAME = config['nacos']['cluster_name']
+    client = nacos.NacosClient(NACOS_SERVER, namespace=NAMESPACE, username=config['nacos']['username'], password=config['nacos']['password'])
+    SERVICE_NAME = config['nacos']['service_name']
+    HEALTH_CHECK_INTERVAL = int(config['nacos']['health_check_interval'])
+    if config.has_option('nacos', 'weight'):
+        WEIGHT = int(config.get('nacos', 'weight'))
+    else:
+        WEIGHT = 1
+    HTTP_PROXY = config.getboolean('server', 'http_proxy')
+    DOMAIN = config['server']['domain']
+    PROTOCOL = config['server']['protocol']
+    DIRECT_DOMAIN = config.getboolean('server', 'direct_domain')
+    # Parse AutoDLServiceURL
+    autodl_url = os.environ.get('AutoDLServiceURL')
+
+    if not autodl_url:
+        raise RuntimeError("Error: AutoDLServiceURL environment variable is not set.")
+
+    parsed_url = urlparse(autodl_url)
+    SERVICE_IP = parsed_url.hostname
+    SERVICE_PORT = parsed_url.port
+    if not SERVICE_IP or not SERVICE_PORT:
+        raise RuntimeError("Error: Invalid AutoDLServiceURL format.")
 
-	if not args.test_load:
-		uvicorn.run(app, host=args.host, port=args.port)
+    print(f"Service will be registered with IP: {SERVICE_IP} and Port: {SERVICE_PORT}")
+    if not register_service(client,SERVICE_NAME,SERVICE_IP,SERVICE_PORT,CLUSTER_NAME,HEALTH_CHECK_INTERVAL,WEIGHT,HTTP_PROXY,DOMAIN,PROTOCOL,DIRECT_DOMAIN):
+        raise RuntimeError("Service is healthy but failed to register.")
+    app.register_model('text-moderations-latest', 'text-moderations-stable')
+    app.register_model('text-moderations-005', 'text-moderations-ifmain')
+    uvicorn.run(app, host="0.0.0.0", port=6006)