Update app.py

app.py

@@ -1,23 +1,20 @@
-tokenizer1 = AutoTokenizer.from_pretrained("Emma0123/fine_tuned_model")
-model1 = AutoModelForSequenceClassification.from_pretrained("Emma0123/fine_tuned_model")
+from transformers import pipeline
 
-tokenizer2 = AutoTokenizer.from_pretrained("jonas/roberta-base-finetuned-sdg")
-model2 = AutoModelForSequenceClassification.from_pretrained("jonas/roberta-base-finetuned-sdg")
+# 初始化两个模型
+pipe1 = pipeline("text-classification", model="Emma0123/fine_tuned_model")
+pipe2 = pipeline("text-classification", model="jonas/roberta-base-finetuned-sdg")
 
-# 输入文本
-input_text = input()
+# 获取用户输入的文本
+input_text = input("Please enter the text: ")
 
-# 对第一个模型进行推理
-inputs = tokenizer1(input_text, return_tensors="pt", truncation=True)
-outputs = model1(**inputs)
-predictions = torch.argmax(outputs.logits, dim=1).item()
+# 使用第一个模型进行预测
+result1 = pipe1(input_text)
 
-# 根据第一个模型的输出进行条件判断
-if predictions == 1:
-    # 使用第二个模型进行判断
-    inputs2 = tokenizer2(input_text, return_tensors="pt", truncation=True)
-    outputs2 = model2(**inputs2)
-    predictions2 = torch.argmax(outputs2.logits, dim=1).item()
-    print("Second model prediction:", predictions2)
+# 判断第一个模型的输出结果
+if result1[0]['label'] == 'LABEL_1':  # 根据您的模型实际返回的标签进行修改
+    # 如果输出结果为1（或者对应的标签），将输入文本传递给第二个模型
+    result2 = pipe2(input_text)
+    print(result2)
 else:
-    print("This content is unrelated to Environment.")
+    # 如果输出结果为0（或者对应的标签），打印提示信息
+    print("This content is unrelated to Environment.")