Update app.py

app.py

@@ -1,3 +1,4 @@
+import spaces
 import gradio as gr
 import requests
 import json
@@ -5,6 +6,10 @@ import time
 import random
 import re
 import os
+from transformers import pipeline
+
+@spaces.GPU
+pipe = pipeline("text-generation", model="jslin09/gemma2-2b-ner", device="cuda")
 
 elements = {'LEO_SOC': ('犯罪主體', 'Subject of Crime'),
             'LEO_VIC': ('客體', 'Victim'),
@@ -130,57 +135,21 @@ def tag_in_color(tag_content, tag):
     return parts
 
 def le_tagger(content, tag):
-#    api_base_url = "http://localhost:11434/api" # 本機端的 ollama
-    api_base_url = os.environ.get("REMOTE_LLM_SERVER")
-    url = f"{api_base_url}/generate"
-    list_url = f"{api_base_url}/tags"
-    headers = {
-        "Content-Type": "application/json"
-    }
+    jaccard_score = [] # 紀錄單篇的分數
+    score = ()
     random.seed(time.localtime(time.time())[5]) # 以時間為基礎的隨機種子
     tag = tag # 傳進來的標籤英文字串，不含方括號。
     prompt = get_prompt(content, tag, elements[tag][0]) # 將要標註的內以及提示詞組合成作業提示詞
-    ner_model_name = 'gemma2-ner:2b' # 已載入 Ollama 中的 LE-NER 模型名稱
-    data = {
-                "model": f"{ner_model_name}",
-                "prompt": f"{prompt}",
-            #    "format": "json",
-                "options": {
-                    "seed": 60,
-                    "top_k": 50,
-                    "top_p": 0.7,
-                    "temperature": 0.9
-                  },
-                "stream": False
-            }
     loop_time = time.time()
-    jaccard_score = [] # 紀錄單篇的分數
-    score = ()
-    response = requests.post(url, headers=headers, data=json.dumps(data))
-    if response.status_code == 200:
-        print(f'構成要件要素標籤: {tag}\n構成要件要素:')
-        response_text = response.text
-        data = json.loads(response_text)
-    #    print(tokenizer.decode(data["context"])) # 全部內容，需要解碼，
-        actual_response = data["response"] # 只有回應的內容，而且是文字形態，不必解碼。
-        print(f"{actual_response}\n") # 輸出含有標籤的回應內容
-        color_result = tag_in_color(actual_response, tag) # 將回應內容進行標籤標色
-        print(color_result) # 輸出有標色的結果
-        score = Jaccard(actual_response, tag) # 算分數
-        jaccard_score.append(score)
-        print(f"--- 執行耗時：{(time.time() - loop_time)}秒 ---\n")
-    #                    example_num  = example_num  + 1
-    else:
-        # 列出 Ollama 中已經安裝的模型名稱
-        print(f"Error:{response.status_code} {json.loads(response.text)['error']}")
-        response = requests.get(list_url, headers=headers, data=json.dumps(data))
-        response_text = response.text
-    #    print(response_text)
-        data = json.loads(response_text)
-        actual_response = data["models"]
-        for _, model in enumerate(data["models"]):
-            print(model['name'])
-        print(f"--- 執行耗時：{(time.time() - loop_time)}秒 ---\n")
+    pipe_result = pipe(prompt, max_new_tokens=1200)
+#    print(f"Pipe的回應內容:\n{pipe_result[0]['generated_text'].split('標註結果:')[1]}")
+    actual_response = pipe_result[0]['generated_text']
+#    print(actual_response)
+    color_result = tag_in_color(actual_response, tag) # 將回應內容進行標籤標色
+#    print(color_result)
+#    score = Jaccard(actual_response, tag) # 算分數
+#    jaccard_score.append(score)
+    print(f"--- 執行耗時：{(time.time() - loop_time)}秒 ---\n")
     return color_result
 
 examples = ["錢旺來雖明知不法犯罪集團經常使用人頭帳戶，向被害人施用詐術，藉此逃避檢警人員之追緝，並預見向其取得帳戶之人，會以該帳戶作為詐欺取財之不法所用，竟仍基於幫助詐欺之犯意，於民國106年1月11日某時，將其所申辦之上海商業銀行（下稱上海商銀）帳號00000000000000號帳戶之提款卡、密碼，以宅急便之方式，寄送予某真實姓名年籍不詳之詐欺集團成員，供該詐欺集團成員用於詐欺取財之犯行。嗣該詐欺集團成員即意圖為自己不法之所有，於106年1月13日18時許，撥打電話予洪菁霞，向洪菁霞佯稱係其友人，急需資金周轉，致洪菁霞陷於錯誤，於同日14時35分許，至台中市○區○○路000號之郵局，臨櫃匯款新臺幣（下同）10萬元至錢旺來前揭上海商銀帳戶內，隨即遭提領一空。",