Update README.md

README.md

@@ -9,82 +9,82 @@ tags:
 
 使用方法如下：
 
+```python
+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可用，将输入数据移动到GPU
-        #inputs = {key: value.to(device) for key, value in inputs.items()}
+    # 如GPU可用，将输入数据移动到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
-
-    # 示例调用
-    text = "胎禽消息渺难知，小萼妆容故故迟。城郭渐随寒碧敛，湖山刚与晚阴宜，再来恐或成孤往，此去何由问所之。坐对空亭喧冻雀，可堪暝色向人垂。"
-    result = classify_text(text)
-    print(result)
+    return result_str
 
+# 示例调用
+text = "胎禽消息渺难知，小萼妆容故故迟。城郭渐随寒碧敛，湖山刚与晚阴宜，再来恐或成孤往，此去何由问所之。坐对空亭喧冻雀，可堪暝色向���垂。"
+result = classify_text(text)
+print(result)
+```