"""Generator for the repositories of Fujitsu-LLM-KG-8x7B models.
"""

from typing import Literal, Optional
from dataclasses import dataclass
import re
import json

@dataclass(kw_only=True)
class ModelContent:
    company_name = "Fujitsu-LLM-KG"
    model_name: str
    base_model_id: str
    train_datasets: str
    how_to_use: str

    @property
    def model_id(self) -> str:
        return f"{self.company_name}/{self.model_name}"


cpt_model = ModelContent(
        model_name="Fujitsu-LLM-KG-8x7B_cpt",
        base_model_id="mistralai/Mixtral-8x7B-Instruct-v0.1",
        train_datasets="""
            ### Continual Pre-Training

            * ナレッジグラフ対訳コーパス（日・英）
                * *cf., [当社の技術ブログ](https://blog.fltech.dev/entry/2024/10/15/Fujitsu-LLM-KG-ja)*
            * ナレッジグラフ関連ソースコード（日・英）
                * *e.g., RDF Turtle、SPARQL、PlantUML、など*
            * ナレッジグラフ関連コーパス（日・英）
                * *"ナレッジグラフ"、"知識トリプル"、などでフィルタリング*
            * Webクローリングデータ（日）
            * Wikipediaデータ（日・英）
            * 論文コーパス（日）
            * 法律データ（日）
            * 判例コーパス（日）
            * 英語数学コーパス（英）
            * 法律対訳コーパス（日英）
            * 字幕対訳コーパス（日英）
            """,
        how_to_use='''
            ### Generates Knowledge Graph from Text

            タスクを指示する。
            ````python
            prompt = """
            [INST]
            Generate "Knowledge Graph" in RDF Turtle format based on the given "Source".

            ## Source
            ```txt
            宗像聡は富士通に2010年から勤めています。
            彼はFujitsu-LLM-KG-8x7B_cptを開発しました。
            ```

            ## Strategy
            Extract all verifiable facts in "Source" as knowledge triples.
            [/INST]
            """.strip()
            
            generated = kgllm.generate(prompt)
            print(generated)
            ````

            結果を確認する。
            ````md
            ## Knowledge Graph
            ```turtle
            #@rationale: 宗像聡は富士通に2010年から勤めています。
            <#宗像聡>
                rel:employer [
                    rdf:object <#富士通>;
                    rel:start_time <#2010>
                ].

            #@rationale: 彼はFujitsu-LLM-KG-8x7B_cptを開発しました。
            <#宗像聡>
                rel:notable_work <#Fujitsu-LLM-KG-8x7B_cpt>.
            ```
            ````

            ### Generates Text from Knowledge Graph

            タスクを指示する。
            ````python
            prompt = """
            Generate "Text" to explain the given knowledge triples in "Source".

            ## Source
            ```turtle
            <#Satoshi Munakata>
                rel:notable_work <#Fujitsu-LLM-KG-8x7B_cpt>;
                rel:employer [
                    rdf:object <#Fujitsu>;
                    rel:start_time <#2010>
                ].
            ```

            ## Strategy
            Explain the knowledge triples in "Source" without omission, but concisely and fluently.
            [/INST]
            """.strip()

            generated = kgllm.generate(prompt)
            print(generated)
            ````

            結果を確認する。
            ````md
            ## Text
            ```txt
            Satoshi Munakata, who started working for Fujitsu in 2010, is the creator of the notable work Fujitsu-LLM-KG-8x7B cpt.
            ```
            ````
            ''',
    )

model_contents = [
    cpt_model,
    ModelContent(
        model_name="Fujitsu-LLM-KG-8x7B_inst-infer_v1",
        base_model_id=cpt_model.model_id,
        train_datasets="""
            ### Instruction Tuning

            * ナレッジグラフ推論用の指示学習データ ver.1
                * [JEMHopQA](https://github.com/aiishii/JEMHopQA)から合成したタスクデータ
            """,
        how_to_use='''
            ### Generates Graph Schema

            *（以下のタスク例は、[JEMHopQA](https://github.com/aiishii/JEMHopQA)を改変したもののため、[CC BY-SA 4.0ライセンス](https://creativecommons.org/licenses/by-sa/4.0/deed.en)が継承されます）*

            タスクを指示する。
            ````python
            question = "長岡京に遷都される前の都の所在地は現在の何県の何市にあたるでしょうか？"

            prompt_schema = f"""
            [INST]
            Generate "Knowledge Graph" in RDF Turtle format based on the given "Source".

            ## Source
            ```txt
            {question}
            ```

            ## Strategy
            Extract graph schema needed to answer the question in above "Source" as knowledge triples without omission.
            [/INST]
            """

            generated_schema = kgllm.generate(prompt_schema)
            print(generated_schema)
            ````

            結果を確認する。
            ````md
            ## Knowledge Graph
            ```turtle
            <#長岡京>
                rel:遷都される前の都 <#?>.

            <#?>
                rel:現在の県 <#?>;
                rel:現在の市 <#?>.
            ```
            ````

            ### Resolves to Knowledge Graph

            *（以下のタスク例は、[Wikipedia](https://ja.wikipedia.org/wiki/)を改変したもののため、[CC BY-SA 3.0ライセンス](https://creativecommons.org/licenses/by-sa/3.0/deed.en)が継承されます）*

            タスクを指示する。
            ````python
            prompt_kg = f"""
            [INST]
            Generate "Knowledge Graph" in RDF Turtle format based on the given "Source".
            
            ## Source
            ```txt
            平安京

            平安京（へいあんきょう／たいらのみやこ）または平安城（へいあんじょう）は、日本における古代最後の宮都。794年（延暦13年）から1869年（明治2年）までの日本の首都。
            桓武天皇により、長岡京に代わる都として山背国（山城国）愛宕・葛野の両郡にまたがる地が選ばれ、中国の洛陽城や長安城を模して793年（延暦12年）から建設された。翌794年（延暦13年）に遷都。北部中央に宮城・平安宮（大内裏）が建設され、以降歴代の皇居が置かれた。
            遷都以来、平清盛により断行された福原遷都（1180年）の期間を除いて、東京奠都まで1100年近くに亘って都として機能し、1869年（明治2年）まで続いた。今日の京都市街が形成されるに至る。
            ----
            平城京

            平城京（へいじょうきょう／へいぜいきょう／ならのみやこ）は、奈良時代の日本の首都。710年に藤原京から遷都するにあたり、唐の都長安城を模倣して大和国に建造された都城。現在の奈良県奈良市、大和郡山市に存在する。
            中央北域に宮城・平城宮（大内裏）を置き、東西8坊 (約 4.3 km) の面積をもち、中央を南北に走る朱雀大路によって左京・右京に二分され、さらに南北・東西を大路・小路によって碁盤の目のように整然と区画され、全域が72坊に区画設定されていた。
            大阪湾や淡路島（八島のひとつ）にも近い奈良盆地（奈良県奈良市の西部の一部、中心部及び大和郡山市北部）には、5世紀中頃にはすでに天皇陵である佐紀盾列古墳群が作られ、またのちには大神神社、7世紀には興福寺も建立されているが、京となった8世紀には、東大寺や巨大な仏像である東大寺盧舎那仏像、法華寺などが建立された。本州の政治・文化の中心地となるに至って外京（げきょう）に位置した門前町が、今に続く奈良の町を形成する中心となった。
            ----
            長岡京

            長岡京（ながおかきょう）は、山城国乙訓郡にあった奈良時代末期（または平安時代初期）の都城（現在の京都府向日市、長岡京市、京都市西京区）。宮域跡は向日市鶏冠井町（かいでちょう）に位置し、「長岡宮跡」として国の史跡に指定されている。
            延暦3年（784年）11月11日、第50代桓武天皇により平城京から遷都され、延暦13年（794年）10月22日に平安京に遷都されるまで機能した。
            ```

            ## Strategy
            Extract all verifiable facts in above "Source", that match the following graph schema, as knowledge triples without omission.
            ```turtle
            {extract_turtle(generated_schema)}
            ```
            The extracted facts are needed to answer the question "{question}".
            [/INST]
            """

            generated_kg = kgllm.generate(prompt_kg)
            print(generated_kg)
            ````

            結果を確認する。
            ````md
            ## Knowledge Graph
            ```turtle
            <#長岡京>
                rel:遷都される前の都 <#平城京>.

            <#平城京>
                rel:現在の県 <#奈良県>;
                rel:現在の市 <#奈良市>.
            ```
            ````

            ### Generate Final Answer

            タスクを指示する。
            ````python
            prompt_answer = f"""
            [INST]
            Explore "Knowledge Graph" entity-to-entity then finally answer "Question".

            ## Knowledge Graph
            ```turtle
            {extract_turtle(generated_kg)}
            ```
            
            ## Question
            {question}
            
            ## Strategy
            Answer briefly in one line.
            [/INST]
            """

            generated_answer = kgllm.generate(prompt_answer)
            print(generated_answer)
            ````

            結果を確認する。
            ````md
            ## Explore Path
            ```path
            長岡京 → rel:遷都される前の都 → 平城京
            平城京 → rel:現在の県 → 奈良県
            平城京 → rel:現在の市 → 奈良市
            ```

            ## Answer
            ```txt
            奈良県奈良市
            ```
            ````
            ''',
    ),
    ModelContent(
        model_name="Fujitsu-LLM-KG-8x7B_inst-infer_v2",
        base_model_id=cpt_model.model_id,
        train_datasets="""
            ### Instruction Tuning

            * ナレッジグラフ推論用の指示学習データ ver.2
                * [HotpotQA](https://hotpotqa.github.io/)から合成したタスクデータ
            """,
        how_to_use='''
            ### Generates Graph Schema

            タスクを指示する。
            ````python
            question = "お正月とお盆とでは、どちらがより長期間の休日ですか?"

            prompt_schema = f"""
            [INST]
            Generate "Knowledge Graph" in RDF Turtle format based on the given "Source".

            ## Source
            ```txt
            {question}
            ```
            
            ## Strategy
            Extract graph schema needed to answer the question in above "Source" as knowledge triples without omission.
            [/INST]
            """.strip()

            generated_schema = kgllm.generate(prompt_schema)
            print(generated_schema)
            ````

            結果を確認する。
            ````md
            ## Knowledge Graph
            ```turtle
            <#お正月>
                rel:has_duration <?duration_of_new_year>.
            <#お盆>
                rel:has_duration <?duration_of_obon>.
            <#Answer>
                rel:is <#お正月>;
                rel:when [
                    <?duration_of_new_year>
                        rel:greater_than <?duration_of_obon>.
                ];
                rel:else <#お盆>.
            ```
            ````

            ### Resolves to Knowledge Graph

            *（以下のタスク例は、厚生労働省が公開している[モデル就業規則](https://www.mhlw.go.jp/stf/seisakunitsuite/bunya/koyou_roudou/roudoukijun/zigyonushi/model/index.html)を抜粋・改変したものです。）*

            タスクを指示する。
            ````python
            prompt_kg = f"""
            [INST]
            Generate "Knowledge Graph" in RDF Turtle format based on the given "Source".

            ## Source
            ```txt
            Title: 第１９条「労働時間及び休憩時間」

            １　労働者代表と１年単位の変形労働時間制に関する労使協定を締結した場合、当該協定の適用を受ける労働者について、１週間の所定労働時間は、対象期間を平均して１週間当たり４０時間とする。
            ２　１年単位の変形労働時間制を適用しない労働者について、１週間の所定労働時間は４０時間、１日の所定労働時間は８時間とする。

            ----
            Title: 第２０条「休日」

            １　１年単位の変形労働時間制の適用を受ける労働者の休日については、１年単位の変形労働時間制に関する労使協定の定めるところにより、対象期間の初日を起算日とする１週間ごとに１日以上、１年間に１２５日以上となるように指定する。その場合、年間休日カレンダーに定め、対象期間の初日の３０日前までに各労働者に通知する。
            ２　１年単位の変形労働時間制を適用しない労働者の休日については、以下のとおり指定し、月間休日カレンダーに定め、対象期間の初日の３０日前までに各労働者に通知する。
            ①日曜日（前条第３号の特定期間を除く。）
            ②国民の祝日（日曜日と重なったときは翌日）
            ③年末年始（１月１日～１月３日）
            ④夏季休日（８月１３日～８月１６日）
            ⑤その他会社が指定する日

            ----
            Title: 第２１条「時間外及び休日労働等」

            １　業務の都合により、第１９条の所定労働時間を超え、又は第２０条の所定休日に労働させることがある。
            ２　前項の場合、法定労働時間を超える労働又は法定休日における労働については、あらかじめ会社は労働者の過半数代表者と書面による労使協定を締結するとともに、これを所轄の労働基準監督署長に届け出るものとする。
            ３　妊娠中の女性、産後１年を経過しない女性労働者（以下「妊産婦」という）であって請求した者及び１８歳未満の者については、第２項による時間外労働又は休日若しくは深夜（午後１０時から午前５時まで）労働に従事させない。
            ４　災害その他避けることのできない事由によって臨時の必要がある場合には、第１項から前項までの制限を超えて、所定労働時間外又は休日に労働させることがある。ただし、この場合であっても、請求のあった妊産婦については、所定労働時間外労働又は休日労働に従事させない。

            ----
            Title: 第２６条「育児時間」

            １　１歳に満たない子を養育する女性労働者から請求があったときは、休憩時間のほか１日について２回、１回について３０分の育児時間を与える。
            ```
            
            ## Strategy
            Extract all verifiable facts in above "Source", that match the following graph schema, as knowledge triples.
            ```turtle
            {extract_turtle(generated_schema)}
            ```
            The extracted facts are needed to answer the question "{question}".
            However, if no useful facts are found in the above "Source", do not output any triples.
            [/INST]
            """.strip()

            generated_kg = kgllm.generate(prompt_kg)
            print(generated_kg)
            ````

            結果を確認する。
            ````md
            ## Knowledge Graph
            ```turtle
            <#お正月>
                #@rationale: Title: 第２０条「休日」 - ③年末年始（１月１日～１月３日）
                rel:has_duration <#3>;
            <#お盆>
                #@rationale: Title: 第２０条「休日」 - ④夏季休日（８月１３日～８月１６日）
                rel:has_duration <#4>.
            <#Answer>
                rel:is <#お正月>;
                rel:when [
                    <#3>
                        rel:greater_than <#4>.
                ];
                rel:else <#お盆>.
            ```
            ````

            ### Generate Final Answer

            タスクを指示する。
            ````python
            prompt_answer = f"""
            [INST]
            Explore "Knowledge Graph" entity-to-entity then finally answer "Question".

            ## Knowledge Graph
            ```turtle
            {extract_turtle(generated_kg)}
            ```

            ## Question
            {question}

            ## Strategy
            Answer briefly in one line.
            [/INST]
            """.strip()

            generated_answer = kgllm.generate(prompt_answer)
            print(generated_answer)
            ````

            結果を確認する。
            ````md
            ## Explore Path
            ```path
            お正月は3日間の休日です.
            お盆は4日間の休日です.
            Answer is 4 because 4 is greater than 3.
            ```

            ## Answer
            ```txt
            お盆
            ```
            ````
            ''',
    ),
    ModelContent(
        model_name="Fujitsu-LLM-KG-8x7B_inst-gen_ja",
        base_model_id=cpt_model.model_id,
        train_datasets="""
            ### Instruction Tuning

            * ナレッジグラフ生成用の指示学習データ 日本語版
                * [JacRED](https://github.com/YoumiMa/JacRED)から合成したタスクデータ
            """,
        how_to_use='''
            ### Generates Knowledge Graph from Sentences

            *（以下のタスク例は、[Wikipedia](https://ja.wikipedia.org/wiki/)を改変したもののため、[CC BY-SA 3.0ライセンス](https://creativecommons.org/licenses/by-sa/3.0/deed.en)が継承されます）*

            タスクを指示する。
            ````python
            prompt = """
            [INST]
            Generate "Knowledge Graph" in RDF Turtle format based on the given "Source".

            ## Source
            ```txt
            s0. <#マリー・テレーズ・ドートリッシュ:Person[0]>(<#1638 年 9 月 10 日:Date[6]>日-<#1683 年 7月 30 日:Date[7]>)は、<#フランス:Location[1]>王<#ルイ 14 世:Person[2]>の王妃。
            s1. 父は<#スペイン:Location[3]>王<#フェリペ 4 世:Person[4]>、母は<#フランス:Location[1]>王<#アンリ 4 世:Person[8]>と<#マリー・ド・メディシス:Person[9]>の娘<#イサベル・デ・ボルボン:Person[5]>。
            s2. <#スペイン:Location[3]>名は<#マリア・テレサ:Person[0]>。
            s3. <#ルイ 14 世:Person[2]>の父<#ルイ 13 世:Person[10]>は<#イサベル:Person[5]>の兄、母<#アンヌ・ドートリッシュ:Person[11]>は<#フェリペ 4 世:Person[4]>の姉であり、<#ルイ 14世:Person[2]>と<#マリー・テレーズ:Person[0]>とは父方・母方双方で従兄妹に当たる。
            s4. <#神聖ローマ:Location[12]>皇帝<#レオポルト 1 世:Person[13]>の皇后<#マルガリータ・テレサ:Person[14]>は異母妹、<#スペイン:Location[3]>・<#ハプスブルク家:Organization[15]>最後の王<#カルロス 2世:Person[16]>は異母弟である。
            ```
            ## Strategy
            Extract all verifiable facts about relation "CountryOfCitizenship" in "Source" as knowledge triples.
            * Don't speculate, extract only the facts as stated in "Source" **WITHOUT OMMISION**. 
            The relation "CountryOfCitizenship" (P27) means "the object is a country that recognizes the subject as its citizen".
            * The existence of a relation is restricted by the type of subject and object. You must select subjects and objects of triples from the bracketed eigenexpressions, such as "<#name:type[id]>", without modification.
            [/INST]
            """.strip()

            generated = kgllm.generate(prompt)
            print(generated)
            ````

            結果を確認する。
            ````md
            ## Knowledge Graph
            ```turtle
            #@rationale: s0. <#マリー・テレーズ・ドートリッシュ:Person[0]>(<#1638 年 9 月 10 日:Date[6]>日-<#1683 年 7月 30 日:Date[7]>)は、<#フランス:Location[1]>王<#ルイ 14 世:Person[2]>の王妃。

            <#ルイ 14 世:Person[2]>
                rel:CountryOfCitizenship <#フランス:Location[1]>.

            #@rationale: s1. 父は<#スペイン:Location[3]>王<#フェリペ 4 世:Person[4]>、母は<#フランス:Location[1]>王<#アンリ 4 世:Person[8]>と<#マリー・ド・メディシス:Person[9]>の娘<#イサベル・デ・ボルボン:Person[5]>。

            <#フェリペ 4 世:Person[4]>
                rel:CountryOfCitizenship <#スペイン:Location[3]>.

            <#アンリ 4 世:Person[8]>
                rel:CountryOfCitizenship <#フランス:Location[1]>.

            <#イサベル・デ・ボルボン:Person[5]>
                rel:CountryOfCitizenship <#フランス:Location[1]>.

            #@rationale: s2. <#スペイン:Location[3]>名は<#マリア・テレサ:Person[0]>。

            #@rationale: s3. <#ルイ 14 世:Person[2]>の父<#ルイ 13 世:Person[10]>は<#イサベル:Person[5]>の兄、母<#アンヌ・ドートリッシュ:Person[11]>は<#フェリペ 4 世:Person[4]>の姉であり、<#ルイ 14世:Person[2]>と<#マリー・テレーズ:Person[0]>とは父方・母方双方で従兄妹に当たる。

            <#ルイ 13 世:Person[10]>
                rel:CountryOfCitizenship <#フランス:Location[1]>.

            <#アンヌ・ドートリッシュ:Person[11]>
                rel:CountryOfCitizenship <#フランス:Location[1]>.

            #@rationale: s4. <#神聖ローマ:Location[12]>皇帝<#レオポルト 1 世:Person[13]>の皇后<#マルガリータ・テレサ:Person[14]>は異母妹、<#スペイン:Location[3]>・<#ハプスブルク家:Organization[15]>最後の王<#カルロス 2世:Person[16]>は異母弟である。

            <#カルロス 2世:Person[16]>
                rel:CountryOfCitizenship <#スペイン:Location[3]>.
            ```
            ````
            ''',
    ),
    ModelContent(
        model_name="Fujitsu-LLM-KG-8x7B_inst-gen_en",
        base_model_id=cpt_model.model_id,
        train_datasets="""
            ### Instruction Tuning

            * ナレッジグラフ生成用の指示学習データ 英語版
                * [Re-DocRED](https://github.com/tonytan48/Re-DocRED)から合成したタスクデータ
            """,
        how_to_use='''
            ### Generates Knowledge Graph from Sentences
            
            *（以下のタスク例は、[Wikipedia](https://en.wikipedia.org/wiki/)を改変したもののため、[CC BY-SA 3.0ライセンス](https://creativecommons.org/licenses/by-sa/3.0/deed.en)が継承されます）*

            タスクを指示する。
            ````python
            prompt = """
            [INST]
            Generate "Knowledge Graph" in RDF Turtle format based on the given "Source".
            
            ## Source
            ```txt
            s0. <#John Gregory Myre:Person[0]> is an <#American:Location[1]> journalist and <#NPR:Organization[2]>'s digital editor for international news.
            s1. Previously, he was a foreign correspondent for The <#Associated Press:Organization[3]> and The <#New York Times:Organization[4]> for <#20 years:Time[5]>.
            s2. He reported from <#Jerusalem:Location[6]> from <#2003:Time[7]> to <#2007:Time[8]> for The <#New York Times:Organization[4]>.
            s3. With the <#AP:Organization[3]>, <#Myre:Person[0]> was based in <#Jerusalem:Location[6]>, <#Moscow:Location[9]>, <#Nicosia:Location[10]>, <#Islamabad:Location[11]> and <#Johannesburg:Location[12]>.
            s4. He covered many major international stories, including <#Nelson Mandela:Person[13]>'s release from prison in <#1990:Time[14]>, the <#Gulf War:Miscellaneous[15]> of <#1991:Time[16]>, the rise of the <#Taliban:Organization[17]>, the early days of <#Vladimir Putin:Person[18]> as <#Russia:Location[19]>'s leader, the <#Israeli:Location[20]>-<#Palestinian:Location[21]> fighting that erupted in <#2000:Time[22]> and the <#U.S.:Location[23]> war in <#Afghanistan:Location[24]>.
            s5. He is also a scholar at the <#Middle East Institute:Organization[25]> in <#Washington:Location[26]>.
            ```
            
            ## Strategy
            Extract all verifiable facts about relation "country_of_citizenship" in "Source" as knowledge triples.
            * Don't speculate, extract only the facts as stated in "Source" **WITHOUT OMMISION**.
            The relation "country_of_citizenship" (P27) means "the object is a country that recognizes the subject as its citizen".
            * The existence of a relation is restricted by the type of subject and object.
            You must select subjects and objects of triples from the bracketed eigenexpressions, such as "<#name:type[id]>", without modification.
            [/INST]
            """.strip()

            generated = kgllm.generate(prompt)
            print(generated)
            ```
            ````

            結果を確認する。
            ````md
            ## Knowledge Graph
            ```turtle
            #@rationale: s0. <#John Gregory Myre:Person[0]> is an <#American:Location[1]> journalist and <#NPR:Organization[2]>'s digital editor for international news.

            <#John Gregory Myre:Person[0]>
                rel:country_of_citizenship <#American:Location[1]>.

            #@rationale: s1. Previously, he was a foreign correspondent for The <#Associated Press:Organization[3]> and The <#New York Times:Organization[4]> for <#20 years:Time[5]>.

            #@rationale: s2. He reported from <#Jerusalem:Location[6]> from <#2003:Time[7]> to <#2007:Time[8]> for The <#New York Times:Organization[4]>.

            #@rationale: s3. With the <#AP:Organization[3]>, <#Myre:Person[0]> was based in <#Jerusalem:Location[6]>, <#Moscow:Location[9]>, <#Nicosia:Location[10]>, <#Islamabad:Location[11]> and <#Johannesburg:Location[12]>.

            #@rationale: s4. He covered many major international stories, including <#Nelson Mandela:Person[13]>'s release from prison in <#1990:Time[14]>, the <#Gulf War:Miscellaneous[15]> of <#1991:Time[16]>, the rise of the <#Taliban:Organization[17]>, the early days of <#Vladimir Putin:Person[18]> as <#Russia:Location[19]>'s leader, the <#Israeli:Location[20]>-<#Palestinian:Location[21]> fighting that erupted in <#2000:Time[22]> and the <#U.S.:Location[23]> war in <#Afghanistan:Location[24]>.

            <#Vladimir Putin:Person[18]>
                rel:country_of_citizenship <#Russia:Location[19]>.

            #@rationale: s5. He is also a scholar at the <#Middle East Institute:Organization[25]> in <#Washington:Location[26]>.
            ```
            ````
            ''',
    ),
]


def generate_document(mc: ModelContent, 
        *, 
        kind: Literal["README.md", "LICENSE"], 
        terms_of_use: Optional[str] = None,
    ) -> str:
    """Generate REDAME.md or LICENSE.
    """
    with open(f"./{kind}.template", "r") as f:
        template = f.read()

    file_text = template.format(**dict(
        model_name=mc.model_name,
        base_model_id=mc.base_model_id,
        train_datasets=re.sub(r'^ {12,12}', '', mc.train_datasets, flags=re.MULTILINE).strip(),
        how_to_use=re.sub(r'^ {12,12}', '', mc.how_to_use, flags=re.MULTILINE).strip(),
        terms_of_use=terms_of_use,
    ))
    file_path = f"../../{mc.model_name}/{kind}"

    with open(file_path, "w") as f:
        f.write(file_text)
        print(f"Output to {file_path}")
    
    return file_text


if __name__ == "__main__":
    for mc in model_contents:
        terms_of_use = generate_document(mc, kind="LICENSE")
        generate_document(mc, kind="README.md", terms_of_use=json.dumps(terms_of_use))