add sonnet support

app.py

@@ -75,13 +75,14 @@ def do( business_id, business_name, address):
     extracted_results = extract_results( crawled_results, classes=classes)
     # logger.error(extracted_results['extracted_results'].columns)
     extracted_results = extracted_results['extracted_results'][ [ 'business_id', 'business_name', 'address', 'category', 'evidence', 'phone_number', 'description', 'store_name'] ]
-
+    logger.debug( extracted_results['category'])
+    print(extracted_results['category'])
     postprocessed_results = postprocess_result( extracted_results, postprocessed_results_path="/tmp/postprocessed_results.joblib", category_hierarchy=category2supercategory)
     os.remove("/tmp/postprocessed_results.joblib")
 
     formatted_results = format_output( postprocessed_results)
-    # logger.error( formatted_results.columns)
-    
+    logger.debug( formatted_results)
+    print(formatted_results)
     formatted_output = format_category( formatted_results)
 
     img = plot_wordcloud(formatted_results['formatted_evidence'].values[0])

model.py

@@ -0,0 +1,130 @@
+import os
+import json
+import argparse
+
+from dotenv import load_dotenv
+import anthropic
+from openai import OpenAI
+
+from utils import parse_json_garbage
+
+load_dotenv()
+
+def llm( provider, model, system_prompt, user_content):
+    """Invoke LLM service
+    Argument
+    --------
+    provider: str
+        openai or anthropic
+    model: str
+        Model name for the API
+    system_prompt: str
+        System prompt for the API
+    user_content: str
+        User prompt for the API
+    Return
+    ------
+    response: str
+    """
+    if provider=='openai':
+        client = OpenAI( organization = os.getenv('ORGANIZATION_ID'))
+        chat_completion = client.chat.completions.create( 
+            messages=[
+                {
+                    "role": "system",
+                    "content": system_prompt
+                },
+                {
+                    "role": "user",
+                    "content": user_content,
+                }
+            ],
+            model = model, 
+            response_format = {"type": "json_object"},
+            temperature = 0,
+            # stream = True
+        )
+        response = chat_completion.choices[0].message.content
+
+    elif provider=='anthropic':
+        client = anthropic.Client(api_key=os.getenv('ANTHROPIC_APIKEY'))
+        response = client.messages.create(
+            model= model,
+            system= system_prompt,
+            messages=[
+                {"role": "user", "content": user_content} # <-- user prompt
+            ],
+            max_tokens = 1024
+        )
+        response = response.content[0].text
+    else:
+        raise Exception("Invalid provider")
+    
+    return response
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--provider", type=str, default='anthropic', help="openai or anthropic")
+    parser.add_argument("--model", type=str, default='claude-3-sonnet-20240229', help="Model name for the API", 
+                        choices = ["claude-3-sonnet-20240229", "claude-3-haiku-20240307", "gpt-3.5-turbo-0125", "gpt-4-0125-preview"])
+    parser.add_argument("--classes", type=list, default=['小吃店', '日式料理(含居酒屋，串燒)', '火(鍋／爐)', '東南亞料理(不含日韓)', '海鮮熱炒',  '特色餐廳(含雞、鵝、牛、羊肉)', '傳統餐廳', '燒烤', '韓式料理(含火鍋，烤肉)', '西餐廳(含美式，義式，墨式)', '西餐廳(餐酒館、酒吧、飛鏢吧、pub、lounge bar)', '西餐廳(土耳其、漢堡、薯條、法式、歐式、印度)', '早餐'])
+    parser.add_argument("--task", type=list, default='extract', choices=['extract', 'classify'])
+    args = parser.parse_args()
+
+
+
+    classes = ['小吃店', '日式料理(含居酒屋，串燒)', '火(鍋／爐)', '東南亞料理(不含日韓)', '海鮮熱炒',  '特色餐廳(含雞、鵝、牛、羊肉)', '傳統餐廳', '燒烤', '韓式料理(含火鍋，烤肉)', '西餐廳(含美式，義式，墨式)', ]
+    backup_classes = [ '中式', '西式']
+    
+    extraction_prompt = '''
+            As a helpful and rigorous retail analyst, given the provided query and a list of search results for the query, 
+            your task is to first identify relevant information of the identical store based on store name and proxmity of address if known. After that, extract `store_name`, `address`, `description`, `category` and `phone_number` from the found relevant information, where `category` can only be `小吃店`, `日式料理(含居酒屋，串燒)`, `火(鍋／爐)`, `東南亞料理(不含日韓)`, `海鮮熱炒`,  `特色餐廳(含雞、鵝、牛、羊肉)`, `傳統餐廳`, `燒烤`, `韓式料理(含火鍋，烤肉)`, `西餐廳(含美式，義式，墨式)`, `西餐廳(餐酒館、酒吧、飛鏢吧、pub、lounge bar)`, `西餐廳(土耳其、漢堡、薯條、法式、歐式、印度)` or `早餐`. 
+            It's very important to omit unrelated results. Do not make up any assumption.
+            Please think step by step, and output in json format. An example output json is like {"store_name": "...", "address": "...", "description": "... products, service or highlights ...", "category": "...", "phone_number": "..."}
+            If no relevant information has been found, simply output json with empty values.
+            I'll tip you and guarantee a place in heaven you do a great job completely according to my instruction.
+        '''
+    classification_prompt = f"""
+            As a helpful and rigorous retail analyst, given the provided information about a store, 
+            your task is two-fold. First, classify provided evidence below into the mostly relevant category from the following: {classes}. 
+            Second, if no relevant information has been found, classify the evidence into the mostly relevant supercategory from the following: {backup_classes}.
+            It's very important to omit unrelated piece of evidence and don't make up any assumption.
+            Please think step by step, and must output in json format. An example output json is like {{"category": "..."}}
+            If no relevant piece of information can ever be found at all, simply output json with empty string "".
+            I'll tip you and guarantee a place in heaven you do a great job completely according to my instruction.
+    """
+    
+    if args.task == 'extract':
+        system_prompt = extraction_prompt
+    elif args.task == 'classify':
+        system_prompt = classification_prompt
+    else:
+        raise Exception("Invalid task")
+    
+    query = "山の迴饗"
+    search_results = str([{"title": "山の迴饗", "snippet": "謝謝大家這麼支持山の迴饗 我們會繼續努力用心做出美味的料理 ————————— ⛰️ 山の迴饗地址：台東縣關山鎮中華路56號訂位專線：0975-957-056 · #山的迴饗 · #夢想起飛"}, {"title": "山的迴饗餐館- 店家介紹", "snippet": "營業登記資料 · 統一編號. 92433454 · 公司狀況. 營業中 · 公司名稱. 山的迴饗餐館 · 公司類型. 獨資 · 資本總額. 30000 · 所在地. 臺東縣關山鎮中福里中華路56號 · 使用發票."}, {"title": "關山漫遊| 💥山の迴饗x night bar", "snippet": "山の迴饗x night bar 即將在12/1號台東關山開幕! 別再煩惱池上、鹿野找不到宵夜餐酒館 各位敬請期待並關注我們✨ night bar❌山的迴饗 12/1 ..."}, {"title": "山的迴饗| 中西複合式餐廳｜焗烤飯｜義大利麵 - 台灣美食網", "snippet": "山的迴饗| 中西複合式餐廳｜焗烤飯｜義大利麵｜台式三杯雞｜滷肉飯｜便當｜CP美食營業時間 ; 星期一, 休息 ; 星期二, 10:00–14:00 16:00–21:00 ; 星期三, 10:00–14:00 16:00– ..."}, {"title": "便當｜CP美食- 山的迴饗| 中西複合式餐廳｜焗烤飯｜義大利麵", "snippet": "餐廳山的迴饗| 中西複合式餐廳｜焗烤飯｜義大利麵｜台式三杯雞｜滷肉飯｜便當｜CP美食google map 導航. 臺東縣關山鎮中華路56號 +886 975 957 056 ..."}, {"title": "山的迴饗餐館", "snippet": "山的迴饗餐館,統編:92433454,地址:臺東縣關山鎮中福里中華路56號,負責人姓名:周偉慈,設立日期:112年11月15日."}, {"title": "山的迴饗餐館", "snippet": "山的迴饗餐館. 資本總額(元), 30,000. 負責人, 周偉慈. 登記地址, 看地圖 臺東縣關山鎮中福里中華路56號 郵遞區號查詢. 設立日期, 2023-11-15. 資料管理 ..."}, {"title": "山的迴饗餐館, 公司統一編號92433454 - 食品業者登錄資料集", "snippet": "公司或商業登記名稱山的迴饗餐館的公司統一編號是92433454, 登錄項目是餐飲場所, 業者地址是台東縣關山鎮中福里中華路56號, 食品業者登錄字號是V-202257990-00001-5."}, {"title": "山的迴饗餐館, 公司統一編號92433454 - 食品業者登錄資料集", "snippet": "公司或商業登記名稱山的迴饗餐館的公司統一編號是92433454, 登錄項目是公司/商業登記, 業者地址是台東縣關山鎮中福里中華路56號, 食品業者登錄字號是V-202257990-00000-4 ..."}, {"title": "山的迴饗餐館", "snippet": "負責人, 周偉慈 ; 登記地址, 台東縣關山鎮中福里中華路56號 ; 公司狀態, 核准設立 「查詢最新營業狀況請至財政部稅務入口網 」 ; 資本額, 30,000元 ; 所在縣市 ..."}, {"title": "山的迴饗 | 關山美食｜焗烤飯｜酒吧｜義大利麵｜台式三杯雞｜滷肉飯｜便當｜CP美食", "顧客評價": "324晚餐餐點豬排簡餐加白醬焗烤等等餐點。\t店家也提供免費的紅茶 綠茶 白開水 多種的調味料自取 總而言之 CP值真的很讚\t空間舒適涼爽，店員服務周到"}, {"title": "類似的店", "snippet": "['中國菜']\t['客家料理']\t['餐廳']\t['熟食店']\t['餐廳']"}, {"telephone_number": "0975 957 056"}])
+
+    # query = "大吃一斤泰國蝦麻辣牛肉爐"
+    # search_results = str([{"title": "大吃一斤泰國蝦麻辣牛肉爐", "snippet": "... 一支、本店特賣價600元免費代料理、 保證、活的!歡迎來電預定0975-147-848大吃一斤活蝦料理店新北市三重區自強路一段222號泰國蝦活蝦現場料理不漲價一斤維持一斤480元."}, {"title": "大吃一斤泰國蝦麻辣牛肉爐", "snippet": "... 一支、本店特賣價600元免費代料理、 保證、活的!歡迎來電預定0975-147-848大吃一斤活蝦料理店新北市三重區自強路一段222號泰國蝦活蝦現場料理不漲價一斤維持一斤480元."}, {"title": "大吃一斤", "snippet": "大吃一斤在foodpanda點的到，更多New Taipei City 推薦美食，線上訂立即送，下載foodpanda APP，20分鐘外送上門！瀏覽菜單和獨家優惠折扣."}, {"title": "大吃一斤(新北板橋店)菜單", "snippet": "大吃一斤(新北板橋店) 在foodpanda點的到，更多New Taipei City 推薦美食，線上訂立即送，下載foodpanda APP，20分鐘外送上門！"}, {"title": "大吃一斤活蝦餐廳- 店家介紹", "snippet": "大吃一斤活蝦餐廳. 資本總額. 200000. 代表人. 李錦鴻. 所在區域. 新北市. 所在地. 新北市三重區自強路1段222號(1樓). 商業類型. 獨資. 異動紀錄. 1111108. 營業狀態為: ..."}, {"title": "新北市| 三重區大吃一斤(泰國蝦牛肉料理店)", "snippet": "大吃一斤(泰國蝦牛肉料理店) 餐廳介紹 ; phone icon 電話, 0975 147 848 ; 營業時間, 星期一17:00–04:00 星期二17:00–04:00 星期三17:00–04:00 星期四17:00– ..."}, {"title": "大吃一斤活蝦餐廳", "snippet": "大吃一斤活蝦餐廳. 負責人姓名, 李錦鴻. 地址, 新北市三重區自強路1段222號(1樓). 現況, 核准設立. 資本額(元), 200,000. 組織類型, 獨資. 登記機關, 新北市政府經濟發展局."}, {"title": "【大吃一斤(泰國蝦牛肉料理店)】網友評價- 新北三重區合菜餐廳", "snippet": "大吃一斤(泰國蝦牛肉料理店) - 網友評論、最新食記(132則) 評分: 4.4分。大吃一斤(泰國蝦牛肉料理店)是位於新北三重區的餐廳，地址: 新北市 ... 生猛活海鮮."}, {"title": "大吃一斤生猛海鮮/活魚料理超值優惠方案", "snippet": "大吃一斤生猛海鮮/活魚料理. 電話：0975-147-848. 地址：新北市三重區自強路一段222號. 營業時間：週一至週日17: ..."}, {"title": "大吃一斤三重店 (泰國蝦料理.平價快炒熱炒.各式海鮮)", "顧客評價": "塔香蛤蜊、胡椒蝦、檸檬蝦、胡椒鳳螺 口味不錯食材新鮮 拍照時蛤蜊已經快被小孩吃光\t蝦子不大，店面不大，魚腥味很重，廁所很多蚊子，連菜裡面也有蚊子🦟，根本吃不下去\t新鮮好吃😋老闆人很Nice 推薦鹽烤蝦以及蒜味奶油蝦👍👍👍"}, {"title": "類似的店", "snippet": "['海鮮']\t['海鮮']\t['海鮮']\t['海鮮']"}, {"telephone_number": "0975 147 848"}])
+
+    if args.provider == "openai":
+        client = OpenAI( organization = os.getenv('ORGANIZATION_ID'))
+        # categories = ", ".join([ "`"+x+"`" for x in args.classes if x!='早餐' ])+ " or " + "`早餐`"
+        user_content = f'''
+            `query`: `{query}`,
+            `search_results`: {search_results}
+        '''
+        resp = llm( args.provider, args.model, system_prompt, user_content)
+        print(f"resp -> {resp}")
+
+
+    elif args.provider == "anthropic":
+        client = anthropic.Client(api_key=os.getenv('ANTHROPIC_APIKEY'))
+        user_content = f'''
+            `query`: `{query}`,
+            `search_results`: {search_results}
+        '''
+        print(f"user_content -> {user_content}")
+        resp = llm( args.provider, args.model, system_prompt, user_content)
+        print(resp)

sheet.py

@@ -14,6 +14,9 @@ import tiktoken
 from openai import OpenAI
 from tqdm import tqdm
 
+from model import llm
+from utils import parse_json_garbage
+
 load_dotenv()
 ORGANIZATION_ID = os.getenv('OPENAI_ORGANIZATION_ID')
 SERP_API_KEY = os.getenv('SERP_APIKEY')
@@ -69,69 +72,45 @@ def get_condensed_result(result):
     # print( condensed_results )
     return condensed_result
 
-
-def compose_analysis( client, query, search_results, classes: list, model: str = 'gpt-3.5-turbo-0125'):
+def compose_extraction( query, search_results, classes: list, provider: str, model: str):
     """
     Argument
         query: str
         search_results: str
+        system_prompt: str
+        classes: list, `小吃店`, `日式料理(含居酒屋，串燒)`, `火(鍋／爐)`, `東南亞料理(不含日韓)`, `海鮮熱炒`,  `特色餐廳(含雞、鵝、牛、羊肉)`, `傳統餐廳`, `燒烤`, `韓式料理(含火鍋，烤肉)`, `西餐廳(含美式，義式，墨式)`, `西餐廳(餐酒館、酒吧、飛鏢吧、pub、lounge bar)`, `西餐廳(土耳其、漢堡、薯條、法式、歐式、印度)` or `早餐`
+        provider: "openai"
         model: "gpt-4-0125-preview" or 'gpt-3.5-turbo-0125'
     Return
         response: str
     """
-    categories = ", ".join([ "`"+x+"`" for x in classes if x!='早餐' ])+ " or " + "`早餐`"
-    # print(f"categoreis: {categories}")
-    system_prompt = '''
-                    As a helpful and rigorous retail analyst, given the provided query and a list of search results for the query, 
-                    your task is to first identify relevant information of the identical store based on store name and proxmity of address if known. After that, extract `store_name`, `address`, `description`, `category` and `phone_number` from the found relevant information, where `category` can only be `小吃店`, `日式料理(含居酒屋，串燒)`, `火(鍋／爐)`, `東南亞料理(不含日韓)`, `海鮮熱炒`,  `特色餐廳(含雞、鵝、牛、羊肉)`, `傳統餐廳`, `燒烤`, `韓式料理(含火鍋，烤肉)`, `西餐廳(含美式，義式，墨式)`, `西餐廳(餐酒館、酒吧、飛鏢吧、pub、lounge bar)`, `西餐廳(土耳其、漢堡、薯條、法式、歐式、印度)` or `早餐`. 
-                    It's very important to omit unrelated results. Do not make up any assumption.
-                    Please think step by step, and output in json format. An example output json is like {"store_name": "...", "address": "...", "description": "... products, service or highlights ...", "category": "...", "phone_number": "..."}
-                    If no relevant information has been found, simply output json with empty values.
-                    I'll tip you and guarantee a place in heaven you do a great job completely according to my instruction.
-                '''
-    # print(f"system prompt = {system_prompt}")
-    chat_completion = client.chat.completions.create( 
-        messages=[
-            {
-                "role": "system",
-                "content": system_prompt
-            },
-            {
-                "role": "user",
-                "content": f'''
-                    `query`: `{query}`,
-                    `search_results`: {search_results}
-                ''',
-            }
-        ],
+    classes = ", ".join([ "`"+x+"`" for x in classes if x!='早餐' ])+ " or " + "`早餐`"
+    system_prompt = f'''
+            As a helpful and rigorous retail analyst, given the provided query and a list of search results for the query, 
+            your task is to first identify relevant information of the identical store based on store name and proxmity of address if known. After that, extract `store_name`, `address`, `description`, `category` and `phone_number` from the found relevant information, where `category` can only be {classes}. 
+            It's very important to omit unrelated results. Do not make up any assumption.
+            Please think step by step, and output in json format. An example output json is like {{"store_name": "...", "address": "...", "description": "... products, service or highlights ...", "category": "...", "phone_number": "..."}}
+            If no relevant information has been found, simply output json with empty values.
+            I'll tip you and guarantee a place in heaven you do a great job completely according to my instruction.
+        '''
+    user_content = f"`query`: `{query}`\n`search_results`: {search_results}"
+    response = llm( 
+        provider = provider, 
         model = model, 
-        response_format = {"type": "json_object"},
-        temperature = 0,
-        # stream = True
+        system_prompt = system_prompt, 
+        user_content = user_content
     )
-#     response = []
-#     for chunk in chat_completion:
-#         text = chunk.choices[0].delta.content or "" 
-#         response.append(text)
-#         print( text, end="")
-#     return "".join(response)
-    response = chat_completion.choices[0].message.content
     return response
 
 
-def compose_classication( 
-        client, 
-        evidence, 
-        classes: list = ['小吃店', '日式料理(含居酒屋，串燒)', '火(鍋／爐)', '東南亞料理(不含日韓)', '海鮮熱炒',  '特色餐廳(含雞、鵝、牛、羊肉)', '傳統餐廳', '燒烤', '韓式料理(含火鍋，烤肉)', '西餐廳(含美式，義式，墨式)', ],
-        backup_classes: list = [ '中式', '西式'],
-        model: str = 'gpt-3.5-turbo-0125'
-    ) -> str:
+def compose_classication(  user_content, classes: list, backup_classes: list, provider: str, model: str) -> str:
     """
     Argument
         client: 
         evidence: str
         classes: list
-        model: 'gpt-3.5-turbo-0125', 'gpt-4-0125-preview'
+        provider: e.g. 'openai'
+        model: e.g. 'gpt-3.5-turbo-0125', 'gpt-4-0125-preview'
     Return
         response: str
     """
@@ -141,44 +120,34 @@ def compose_classication(
         pass
     else:
         raise Exception(f"Incorrect classes type: {type(classes)}")
-    chat_completion = client.chat.completions.create( 
-        messages=[
-            {
-                "role": "system",
-                "content": f'''
-                    As a helpful and rigorous retail analyst, given the provided information about a store, 
-                    your task is two-fold. First, classify provided evidence below into the mostly relevant category from the following: {classes}. 
-                    Second, if no relevant information has been found, classify the evidence into the mostly relevant supercategory from the following: {backup_classes}.
-                    It's very important to omit unrelated piece of evidence and don't make up any assumption.
-                    Please think step by step, and output in json format. An example output json is like {{"category": "..."}}
-                    If no relevant piece of information can ever be found at all, simply output json with empty string "".
-                    I'll tip you and guarantee a place in heaven you do a great job completely according to my instruction.
-                '''
-            },
-            {
-                "role": "user",
-                "content": f'''
-                    `evidence`: `{evidence}`
-                ''',
-            }
-        ],
+    system_prompt = f"""
+            As a helpful and rigorous retail analyst, given the provided information about a store, 
+            your task is two-fold. First, classify provided evidence below into the mostly relevant category from the following: {classes}. 
+            Second, if no relevant information has been found, classify the evidence into the mostly relevant supercategory from the following: {backup_classes}.
+            It's very important to omit unrelated piece of evidence and don't make up any assumption.
+            Please think step by step, and must output in json format. An example output json is like {{"category": "..."}}
+            If no relevant piece of information can ever be found at all, simply output json with empty string "".
+            I'll tip you and guarantee a place in heaven you do a great job completely according to my instruction.
+    """
+    response = llm( 
+        provider = provider, 
         model = model, 
-        response_format = {"type": "json_object"},
-        temperature = 0,
-        # stream = True
+        system_prompt = system_prompt,
+        user_content = user_content,
     )
-    response = chat_completion.choices[0].message.content
     return response
 
 
 def classify_results( 
         analysis_results: pd.DataFrame, 
+        classes: list,
+        backup_classes: list,
+        provider: str, 
+        model: str,
         input_column: str = 'evidence', 
         output_column: str = 'classified_category',
-        classes: list = ['小吃店', '日式料理(含居酒屋，串燒)', '火(鍋／爐)', '東南亞料理(不含日韓)', '海鮮熱炒',  '特色餐廳(含雞、鵝、牛、羊肉)', '傳統餐廳', '燒烤', '韓式料理(含火鍋，烤肉)', '西餐廳(含美式，義式，墨式)'],
-        backup_classes: list = [ '中式', '西式']
     ):
-    """
+    """Classify the results
     Argument
         analysis_results: dataframe
         input_column: str
@@ -187,13 +156,13 @@ def classify_results(
     Return 
         analysis_results: dataframe
     """
-    client = OpenAI( organization = ORGANIZATION_ID)
     classified_results = analysis_results.copy()
-    empty_indices = []
-    labels = []
+    labels, empty_indices = [], []
     for idx, evidence in zip( analysis_results['index'], analysis_results[input_column]):
         try:
-            label = json.loads(compose_classication( client, evidence, classes=classes, backup_classes=backup_classes))['category']
+            user_content = f'''`evidence`: `{evidence}`''' 
+            pred_cls = compose_classication( user_content, classes=classes, backup_classes=backup_classes, provider=provider, model=model)
+            label = parse_json_garbage(pred_cls)['category']
             labels.append(label)
         except Exception as e:
             print(f"# CLASSIFICATION error -> evidence: {e}")
@@ -206,13 +175,15 @@ def classify_results(
         "empty_indices": empty_indices
     }
             
-def classify_results_mp( extracted_results: pd.DataFrame, classified_file_path, classes, backup_classes, n_processes: int = 4):
+def classify_results_mp( extracted_results: pd.DataFrame, classified_file_path: str, classes: list, backup_classes: list, provider: str, model: str, n_processes: int = 4):
     """
     Argument
         extracted_results:
         classified_file_path:
-        classes: ['小吃店', '日式料理(含居酒屋，串燒)', '火(鍋／爐)', '東南亞料理(不含日韓)', '海鮮熱炒',  '特色餐廳(含雞、鵝、牛、羊肉)', '傳統餐廳', '燒烤', '韓式料理(含火鍋，烤肉)', '西餐廳(含美式，義式，墨式)']
-        backup_classes: [ '中式', '西式']
+        classes: e.g. ['小吃店', '日式料理(含居酒屋，串燒)', '火(鍋／爐)', '東南亞料理(不含日韓)', '海鮮熱炒',  '特色餐廳(含雞、鵝、牛、羊肉)', '傳統餐廳', '燒烤', '韓式料理(含火鍋，烤肉)', '西餐廳(含美式，義式，墨式)']
+        backup_classes: e.g. [ '中式', '西式']
+        provider:
+        model:
         n_processes: int
     Return
         classified_results: dataframe
@@ -228,10 +199,9 @@ def classify_results_mp( extracted_results: pd.DataFrame, classified_file_path,
                 classify_results,
                 [ ( 
                     d, 
-                    'evidence', 
-                    'classified_category', 
-                    classes,
-                    backup_classes
+                    classes, backup_classes,
+                    provider, model,
+                    'evidence', 'classified_category',
                 ) for d in split_data]
             ) 
             classified_results = merge_results( classified_results, dataframe_columns=['classified_results'], list_columns=['empty_indices'])
@@ -329,16 +299,14 @@ def crawl_results_mp( data: pd.DataFrame, crawl_file_path: str, n_processes: int
     print( f"total time: {time.time() - st}")
     return crawled_results
 
-def extract_results( data: pd.DataFrame, classes: list ):
+def extract_results( data: pd.DataFrame, classes: list, provider: str, model: str):
     """
     Argument
         data: `evidence`, `result`
     Return
         extracted_results: dataframe of `extracted_evidence`
     """
-    client = OpenAI( organization = ORGANIZATION_ID)
-    extracted_results = []
-    empty_indices = []
+    extracted_results, empty_indices, ext_res = [], [], []
     for i, d in tqdm(enumerate(data.itertuples())):
         idx = d[1]
         evidence = d.evidence
@@ -348,10 +316,10 @@ def extract_results( data: pd.DataFrame, classes: list ):
         ana_res = None
         query = compose_query( address, business_name)
         try:
-            ana_res = compose_analysis( client, query = query, search_results = evidence, classes = classes)
-            ana_res = json.loads(ana_res)
+            ext_res = compose_extraction( query = query, search_results = evidence, classes = classes, provider = provider, model = model)
+            ext_res = parse_json_garbage(ext_res)
         except Exception as e:
-            print(f"# ANALYSIS error {e}: i = {i}, ana_res = {ana_res}")
+            print(f"# ANALYSIS error: e = {e}, i = {i}, q = {query}, ext_res = {ext_res}")
             empty_indices.append(i)
             continue
         
@@ -360,7 +328,7 @@ def extract_results( data: pd.DataFrame, classes: list ):
             "business_id": business_id, 
             "business_name": business_name, 
             "evidence": evidence, 
-            ** ana_res
+            ** ext_res
         } )
     extracted_results = pd.DataFrame(extracted_results)
 
@@ -369,9 +337,12 @@ def extract_results( data: pd.DataFrame, classes: list ):
         "empty_indices": empty_indices
     }
 
-def extract_results_mp( crawled_results, extracted_file_path, classes: list):
+def extract_results_mp( crawled_results, extracted_file_path, classes: list, provider: str, model: str, n_processes: int = 4):
     """
     Argument
+        crawled_results: dataframe
+        extracted_file_path
+        classes: list
     Return 
     Reference
         200 records, 4 processes, 502.26914715766907
@@ -380,8 +351,8 @@ def extract_results_mp( crawled_results, extracted_file_path, classes: list):
     # args.extracted_file_path = "data/extracted_results.joblib"
     if not os.path.exists(extracted_file_path):
         split_data = split_dataframe( crawled_results)
-        with mp.Pool(args.n_processes) as pool:
-            extracted_results = pool.starmap( extract_results, [ (x, classes) for x in split_data])
+        with mp.Pool(n_processes) as pool:
+            extracted_results = pool.starmap( extract_results, [ (x, classes, provider, model) for x in split_data])
             extracted_results = merge_results( extracted_results, dataframe_columns=['extracted_results'], list_columns=['empty_indices'])
             with open( extracted_file_path, "wb") as f:
                 joblib.dump( extracted_results, f)
@@ -522,57 +493,59 @@ def main(args):
     Argument
         args: argparse
     """
-    
+    crawled_file_path = os.path.join( args.output_dir, args.crawled_file_path)
+    extracted_file_path = os.path.join( args.output_dir, args.extracted_file_path)
+    classified_file_path = os.path.join( args.output_dir, args.classified_file_path)
+    combined_file_path = os.path.join( args.output_dir, args.combined_file_path)
+    postprocessed_results = os.path.join( args.output_dir, args.postprocessed_results)
+    formatted_results_path = os.path.join( args.output_dir, args.formatted_results_path)
+
     ## 讀取資料名單 ##
-    data = get_leads(args.data_path).tail(20)
+    data = get_leads(args.data_path).tail(5)
     
     ## 進行爬蟲與分析 ##
-    # crawled_results = crawl_results(data)
-    crawled_results = crawl_results_mp( data, args.crawled_file_path, n_processes=args.n_processes)
+    crawled_results = crawl_results_mp( data, crawled_file_path, n_processes=args.n_processes)
+    crawled_results = { k:v[-5:] for k,v in crawled_results.items()}
 
     ## 方法 1: 擷取關鍵資訊與分類 ##
-    # extracted_results = extract_results( 
-    #     crawled_results['crawled_results']
-    # )
     extracted_results = extract_results_mp( 
         crawled_results = crawled_results['crawled_results'], 
-        extracted_file_path = args.extracted_file_path
+        extracted_file_path = extracted_file_path,
+        classes = args.classes,
+        provider = args.provider,
+        model = args.model,
+        n_processes = args.n_processes
     )
 
     ## 方法2: 直接對爬蟲結果分類 ##
-    # classified_results = classify_results(  
-    #     extracted_results['extracted_results'], 
-    #     input_column = 'evidence',
-    #     output_column = 'classified_category',
-    #     classes = ['中式', '西式'],
-    #     backup_classes = [ '中式', '西式']
-    # )
     classified_results = classify_results_mp( 
         extracted_results['extracted_results'], 
-        args.classified_file_path, 
-        classes=args.classes, 
-        backup_classes=args.backup_classes, 
-        n_processes=args.n_processes
+        classified_file_path, 
+        classes = args.classes, 
+        backup_classes = args.backup_classes, 
+        provider = args.provider,
+        model = args.model,
+        n_processes = args.n_processes
     )
 
     ## 合併分析結果 ##
     combined_results = combine_results( 
         classified_results['classified_results'], 
-        args.combined_file_path, 
-        src_column='classified_category', 
-        tgt_column='category', 
-        strategy='replace'
+        combined_file_path, 
+        src_column = 'classified_category', 
+        tgt_column = 'category', 
+        strategy = args.strategy
     ) 
 
     ## 後處理分析結果 ##
     postprossed_results = postprocess_result( 
         combined_results, 
-        args.postprocessed_results,
+        postprocessed_results,
         category2supercategory
     )
 
     formatted_results = format_output( postprossed_results, input_column = 'evidence', output_column = 'formatted_evidence', format_func = format_evidence)
-    formatted_results.to_csv( args.formatted_results, index=False)
+    formatted_results.to_csv( formatted_results_path, index=False)
 
 
 category2supercategory = {
@@ -623,15 +596,18 @@ if __name__=='__main__':
 
     parser = argparse.ArgumentParser()
     parser.add_argument("--data_path", type=str, default="data/餐廳類型分類.xlsx - 測試清單.csv")
-    parser.add_argument("--classified_file_path", type=str, default="data/gpt3.5/classified_results.joblib")
-    parser.add_argument("--extracted_file_path", type=str, default="data/gpt3.5/extracted_results.joblib")
-    parser.add_argument("--crawled_file_path", type=str, default="data/gpt3.5/crawled_results.joblib")
-    parser.add_argument("--combined_file_path", type=str, default="data/gpt3.5/combined_results.joblib")
-    parser.add_argument("--postprocessed_results", type=str, default="data/gpt3.5/postprocessed_results.joblib")
-    parser.add_argument("--formatted_results", type=str, default="data/gpt3.5/formatted_results.csv")
+    parser.add_argument("--output_dir", type=str, help='output directory')
+    parser.add_argument("--classified_file_path", type=str, default="classified_results.joblib")
+    parser.add_argument("--extracted_file_path", type=str, default="extracted_results.joblib")
+    parser.add_argument("--crawled_file_path", type=str, default="crawled_results.joblib")
+    parser.add_argument("--combined_file_path", type=str, default="combined_results.joblib")
+    parser.add_argument("--postprocessed_results", type=str, default="postprocessed_results.joblib")
+    parser.add_argument("--formatted_results_path", type=str, default="formatted_results.csv")
     parser.add_argument("--classes", type=list, default=['小吃店', '日式料理(含居酒屋，串燒)', '火(鍋／爐)', '東南亞料理(不含日韓)', '海鮮熱炒',  '特��餐廳(含雞、鵝、牛、羊肉)', '傳統餐廳', '燒烤', '韓式料理(含火鍋，烤肉)', '西餐廳(含美式，義式，墨式)', '西餐廳(餐酒館、酒吧、飛鏢吧、pub、lounge bar)', '西餐廳(土耳其、漢堡、薯條、法式、歐式、印度)', '早餐'])
     parser.add_argument("--backup_classes", type=list, default=['中式', '西式'])
-    parser.add_argument("--strategy", type=str, default='replace', choices=['replace', 'patch'])
+    parser.add_argument("--strategy", type=str, default='patch', choices=['replace', 'patch'])
+    parser.add_argument("--provider", type=str, default='anthropic', choices=['openai', 'anthropic'])
+    parser.add_argument("--model", type=str, default='claude-3-sonnet-20240229', choices=['claude-3-sonnet-20240229', 'claude-3-haiku-20240307', 'gpt-3.5-turbo-0125', 'gpt-4-0125-preview'])
     parser.add_argument("--n_processes", type=int, default=4)
     args = parser.parse_args()
 

utils.py

@@ -0,0 +1,9 @@
+import json
+
+def parse_json_garbage(s):
+    s = s[next(idx for idx, c in enumerate(s) if c in "{["):]
+    try:
+        return json.loads(s)
+    except json.JSONDecodeError as e:
+        return json.loads(s[:e.pos])
+