Update README.md

README.md

@@ -10,81 +10,105 @@ tags:
 使用方法如下：
 
 
-import torch
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
-import json
-import torch.nn.functional as F
-from zhconv import convert
-import re
+    import torch
+    from transformers import AutoTokenizer, AutoModelForSequenceClassification
+    import json
+    import torch.nn.functional as F
+    from zhconv import convert
+    import re
 
-model_path = "qixun/qilv_classify"
+    model_path = "qixun/qilv_classify"
 
-# 加载模型和分词器
-tokenizer = AutoTokenizer.from_pretrained(model_path)
-model = AutoModelForSequenceClassification.from_pretrained(model_path)
+    # 加载模型和分词器
+    tokenizer = AutoTokenizer.from_pretrained(model_path)
+    model = AutoModelForSequenceClassification.from_pretrained(model_path)
 
-# 如果GPU可用，将模型移动到GPU
-#device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-#model.to(device)
+    # 如果GPU可用，将模型移动到GPU
+    #device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    #model.to(device)
 
-# 加载标签映射关系，label_mapping.json需要根据本机情况修改
-with open("label_mapping.json", "r", encoding="utf-8") as f:
-    label_mapping = json.load(f)
+    # 加载标签映射关系，label_mapping.json需要根据本机情况修改
+    with open("label_mapping.json", "r", encoding="utf-8") as f:
+        label_mapping = json.load(f)
 
 
-def classify_text(text):
+    def classify_text(text):
 
-    text = convert(text, 'zh-cn')
-     # 去掉空格和换行
-    text = text.replace(" ", "").replace("\n", "")
+        text = convert(text, 'zh-cn')
+         # 去掉空格和换行
+        text = text.replace(" ", "").replace("\n", "")
 
-    # 检查文本长度是否为56个字符
-    if len(text) != 64:
-        return "请输入一首带标点的七言律诗"
+        # 检查文本长度是否为56个字符
+        if len(text) != 64:
+            return "请输入一首带标点的七言律诗"
     
-    unique_characters = set(re.findall(r'[\u4e00-\u9fff]', text))
-    if len(unique_characters) < 30:
-        return "请输入一首正常的七言律诗"
+        unique_characters = set(re.findall(r'[\u4e00-\u9fff]', text))
+        if len(unique_characters) < 30:
+            return "请输入一首正常的七言律诗"
     
-    # 准备输入数据
-    inputs = tokenizer(text, padding=True, truncation=True, return_tensors="pt", max_length=512)
+        # 准备输入数据
+        inputs = tokenizer(text, padding=True, truncation=True, return_tensors="pt", max_length=512)
 
-    # 将输入数据移动到GPU
-    inputs = {key: value.to(device) for key, value in inputs.items()}
+        # 将输入数据移动到GPU
+        inputs = {key: value.to(device) for key, value in inputs.items()}
 
-    # 模型推断
-    with torch.no_grad():
-        outputs = model(**inputs)
+        # 模型推断
+        with torch.no_grad():
+            outputs = model(**inputs)
 
-    # 获取预测结果
-    logits = outputs.logits
+        # 获取预测结果
+        logits = outputs.logits
 
-    # 计算每个类别的概率
-    probabilities = F.softmax(logits, dim=-1)
+        # 计算每个类别的概率
+        probabilities = F.softmax(logits, dim=-1)
 
-    # 获取概率最高的三个分类及其概率
-    top_k = 3
-    top_probs, top_indices = torch.topk(probabilities, top_k, dim=-1)
+        # 获取概率最高的三个分类及其概率
+        top_k = 3
+        top_probs, top_indices = torch.topk(probabilities, top_k, dim=-1)
 
-    # 将预测结果转换为标签并附上概率
-    results = []
-    for j in range(top_k):
-        label = label_mapping[str(top_indices[0][j].item())]
-        prob = top_probs[0][j].item()
-        results.append((label, prob))
+        # 将预测结果转换为标签并附上概率
+        results = []
+        for j in range(top_k):
+            label = label_mapping[str(top_indices[0][j].item())]
+            prob = top_probs[0][j].item()
+            results.append((label, prob))
 
-    # 将结果格式化为字符串
-    result_str = "文本: {}\n".format(text)
-    for label, prob in results:
-        result_str += "分类: {}, 概率: {:.4f}\n".format(label, prob)
+        # 将结果格式化为字符串
+        result_str = "文本: {}\n".format(text)
+        for label, prob in results:
+            result_str += "分类: {}, 概率: {:.4f}\n".format(label, prob)
     
-    return result_str
+        return result_str
 
-# 示例调用
-text = "胎禽消息渺难知，小萼妆容故故迟。城郭渐随寒碧敛，湖山刚与晚阴宜，再来恐或成孤往，此去何由问所之。坐对空亭喧冻雀，可堪暝色向人垂。"
-result = classify_text(text)
-print(result)
+    # 示例调用
+    text = "胎禽消息渺难知，小萼妆容故故迟。城郭渐随寒碧敛，湖山刚与晚阴宜，再来恐或成孤往，此去何由问所之。坐对空亭喧冻雀，可堪暝色向人垂。"
+    result = classify_text(text)
+    print(result)
 
 
 
 
+
+
+也可以直接在huggingface里输入一首加标点为64字符的简体七言律诗进行测试，label_mapping.json内容为：
+
+    {
+        "0": "中唐",
+        "1": "乱码",
+        "2": "冲塔",
+        "3": "同光",
+        "4": "复兴",
+        "5": "实验",
+        "6": "晚唐",
+        "7": "江西",
+        "8": "浙",
+        "9": "浣花",
+        "10": "理学",
+        "11": "盛唐",
+        "12": "艳体",
+        "13": "诗界xx",
+        "14": "赣",
+        "15": "闽"
+    }
+
+大家自行转换。