support customize language of tldr

TideDra · Jan 30, 2025 · 331a1b6 · 331a1b6
1 parent 1627409
commit 331a1b6
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 7 deletions.
diff --git a/llm.py b/llm.py
@@ -5,7 +5,7 @@
 GLOBAL_LLM = None
 
 class LLM:
-    def __init__(self, api_key: str = None, base_url: str = None, model: str = None):
+    def __init__(self, api_key: str = None, base_url: str = None, model: str = None,lang: str = "English"):
         if api_key:
             self.llm = OpenAI(api_key=api_key, base_url=base_url)
         else:
@@ -17,6 +17,7 @@ def __init__(self, api_key: str = None, base_url: str = None, model: str = None)
                 verbose=False,
             )
         self.model = model
+        self.lang = lang
 
     def generate(self, messages: list[dict]) -> str:
         if isinstance(self.llm, OpenAI):
@@ -26,9 +27,9 @@ def generate(self, messages: list[dict]) -> str:
             response = self.llm.create_chat_completion(messages=messages,temperature=0)
             return response["choices"][0]["message"]["content"]
 
-def set_global_llm(api_key: str = None, base_url: str = None, model: str = None):
+def set_global_llm(api_key: str = None, base_url: str = None, model: str = None, lang: str = "English"):
     global GLOBAL_LLM
-    GLOBAL_LLM = LLM(api_key=api_key, base_url=base_url, model=model)
+    GLOBAL_LLM = LLM(api_key=api_key, base_url=base_url, model=model, lang=lang)
 
 def get_llm() -> LLM:
     if GLOBAL_LLM is None:

diff --git a/main.py b/main.py
@@ -135,6 +135,12 @@ def get_env(key:str,default=None):
         help="LLM Model Name",
         default="gpt-4o",
     )
+    add_argument(
+        "--language",
+        type=str,
+        help="Language of TLDR",
+        default="English",
+    )
     parser.add_argument('--debug', action='store_true', help='Debug mode')
     args = parser.parse_args()
     assert (
@@ -168,10 +174,10 @@ def get_env(key:str,default=None):
             papers = papers[:args.max_paper_num]
         if args.use_llm_api:
             logger.info("Using OpenAI API as global LLM.")
-            set_global_llm(api_key=args.openai_api_key, base_url=args.openai_api_base, model=args.model_name)
+            set_global_llm(api_key=args.openai_api_key, base_url=args.openai_api_base, model=args.model_name, lang=args.language)
         else:
             logger.info("Using Local LLM as global LLM.")
-            set_global_llm()
+            set_global_llm(lang=args.language)
 
     html = render_email(papers)
     logger.info("Sending email...")

diff --git a/paper.py b/paper.py
@@ -154,13 +154,15 @@ def tldr(self) -> str:
             match = re.search(r'\\section\{Conclusion\}.*?(\\section|\\end\{document\}|\\bibliography|\\appendix|$)', content, flags=re.DOTALL)
             if match:
                 conclusion = match.group(0)
-        prompt = """Given the title, abstract, introduction and the conclusion (if any) of a paper in latex format, generate a one-sentence TLDR summary:
+        llm = get_llm()
+        prompt = """Given the title, abstract, introduction and the conclusion (if any) of a paper in latex format, generate a one-sentence TLDR summary in __LANG__:
         
         \\title{__TITLE__}
         \\begin{abstract}__ABSTRACT__\\end{abstract}
         __INTRODUCTION__
         __CONCLUSION__
         """
+        prompt = prompt.replace('__LANG__', llm.lang)
         prompt = prompt.replace('__TITLE__', self.title)
         prompt = prompt.replace('__ABSTRACT__', self.summary)
         prompt = prompt.replace('__INTRODUCTION__', introduction)
@@ -171,7 +173,7 @@ def tldr(self) -> str:
         prompt_tokens = enc.encode(prompt)
         prompt_tokens = prompt_tokens[:4000]  # truncate to 4000 tokens
         prompt = enc.decode(prompt_tokens)
-        llm = get_llm()
+
         tldr = llm.generate(
             messages=[
                 {