Update generate_markdown.py

generate_markdown.py

@@ -1,95 +1,95 @@
-from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain_community.document_loaders import PyPDFLoader
-from llama_cpp import Llama
-import llama_cpp
-
-def load_llm_model():
-    llm = Llama(
-        model_path="Llama-3.2-3B-Instruct-Q8_0.gguf",
-        # n_gpu_layers = 20, # Uncomment for GPU
-        n_ctx=200000,
-        n_threads=16,
-        n_batch=512,
-        split_mode=llama_cpp.LLAMA_SPLIT_MODE_LAYER,
-        pooling_type=llama_cpp.LLAMA_POOLING_TYPE_RANK,
-        rope_scaling_type=llama_cpp.LLAMA_ROPE_SCALING_TYPE_LINEAR,
-        # main_gpu=0 # Uncomment for GPU
-    )
-    return llm
-
-def get_text_from_pdf(file):
-    loader = PyPDFLoader(file)
-    pages = loader.load_and_split()
-    text_splitter = RecursiveCharacterTextSplitter(chunk_size=250, chunk_overlap=50)
-    texts = text_splitter.split_documents(pages)
-    final_text = ""
-    for text in texts:
-        final_text = final_text + text.page_content
-    print(f"Length of final text: {len(final_text)}")
-    with open("final_text.txt", "w") as f:
-        f.write(final_text)
-    research_paper = ""
-    for line in final_text.split("\n"):
-        if line.startswith("REFERENCES"):
-            break
-        else:
-            research_paper = research_paper + line + " " 
-    with open("research_paper.txt", "w") as f:
-        f.write(research_paper)
-    print(f"Length of research paper: {len(research_paper)}")
-    return research_paper
-
-def generate_prompt(final_text):
-    prompt = f'''
-    You have been provided with a research paper in text format. Your task is to generate a mindmap structure in markdown format that summarizes the research paper.
-    Your output should use the language \"en\" 0.3 times the length of the original research paper. Do not include anything in the response, that is not the part of mindmap and use the following template (any node in the mindmap should not exceed 10-12 words, also generate additional headings that aren't present in document if required for elaborative explaination):
-        # {{Title}} (should be the title of the research paper)
-        ## {{Subtitle01}} (as required and as many as required in markdown format)
-        - {{Emoji01}} Bulletpoint01 (as required and as many as required in markdown format)
-            - {{Emoji01.1}} Bulletpoint01.1 (as required and as many as sub levels required in markdown format)
-                - {{Emoji01.1.1}} Bulletpoint01.1.1 (as required and as many as sub levels required in markdown format)
-                - {{Emoji01.1.2}} Bulletpoint01.1.2 (as required and as many as sub levels required in markdown format)
-            - {{Emoji01.2}} Bulletpoint01.2 (as required and as many as sub levels required in markdown format)
-        - {{Emoji02}} Bulletpoint02 (as required and as many as required in markdown format)
-            - {{Emoji02.1}} Bulletpoint02.1 (as required and as many as sub levels required in markdown format)
-            - {{Emoji02.2}} Bulletpoint02.2 (as required and as many as sub levels required in markdown format)
-                - {{Emoji02.2.1}} Bulletpoint02.2.1 (as required and as many as sub levels required in markdown format)
-                - {{Emoji02.2.2}} Bulletpoint02.2.2 (as required and as many as sub levels required in markdown format)
-                - {{Emoji02.2.3}} Bulletpoint02.2.3 (as required and as many as sub levels required in markdown format)
-                - {{Emoji02.2.4}} Bulletpoint02.2.4 (as required and as many as sub levels required in markdown format)
-        ## {{Subtitle02}} (as required and as many as required in markdown format)
-        - {{Emoji03}} Bulletpoint03 (as required and as many as required in markdown format)
-            - {{Emoji03.1}} Bulletpoint03.1 (as required and as many as sub levels required in markdown format)
-            - {{Emoji03.2}} Bulletpoint03.2 (as required and as many as sub levels required in markdown format)
-                - {{Emoji03.2.1}} Bulletpoint03.2.1 (as required and as many as sub levels required in markdown format)
-                - {{Emoji03.2.2}} Bulletpoint03.2.2 (as required and as many as sub levels required in markdown format)
-        - {{Emoji04}} Bulletpoint04 (as required and as many as required in markdown format)
-            - {{Emoji04.1}} Bulletpoint04.1 (as required and as many as sub levels required in markdown format)
-                - {{Emoji04.1.1}} Bulletpoint04.1.1 (as required and as many as sub levels required in markdown format)
-                - {{Emoji04.1.2}} Bulletpoint04.1.2 (as required and as many as sub levels required in markdown format)
-            - {{Emoji04.2}} Bulletpoint04.2 (as required and as many as sub levels required in markdown format)
-                - {{Emoji04.2.1}} Bulletpoint04.2.1 (as required and as many as sub levels required in markdown format)
-                - {{Emoji04.2.2}} Bulletpoint04.2.2 (as required and as many as sub levels required in markdown format)
-        Summarize the text \"{final_text}\" to generate a elaborated hierarchical mindmap structure (any node in the mindmap should not exceed 10-12 words, also generate additional headings that aren't present in document if required for elaborative explaination) markdown using the \"en\" language 0.3 times the length of the original research paper. Do not include anything in the response, that is not the part of mindmap
-    '''
-    return prompt
-
-def generate_mindmap_structure(llm, prompt):
-    response = llm.create_chat_completion(
-        messages = [
-            {'role':'user',
-            'content': prompt}
-        ],
-        temperature=0.7,
-        top_k=200,
-        top_p=3.0,
-    )
-    mindmap_data = response['choices'][0]['message']['content']
-    print(mindmap_data)
-    return mindmap_data
-
-def generate_markdown(llm, file):
-    final_text = get_text_from_pdf(file)
-    prompt = generate_prompt(final_text)
-    mindmap_markdown = generate_mindmap_structure(llm, prompt)
-    return mindmap_markdown
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_community.document_loaders import PyPDFLoader
+from llama_cpp import Llama
+import llama_cpp
+
+def load_llm_model():
+    llm = Llama(
+        model_path="model.gguf",
+        # n_gpu_layers = 20, # Uncomment for GPU
+        n_ctx=200000,
+        n_threads=16,
+        n_batch=512,
+        split_mode=llama_cpp.LLAMA_SPLIT_MODE_LAYER,
+        pooling_type=llama_cpp.LLAMA_POOLING_TYPE_RANK,
+        rope_scaling_type=llama_cpp.LLAMA_ROPE_SCALING_TYPE_LINEAR,
+        # main_gpu=0 # Uncomment for GPU
+    )
+    return llm
+
+def get_text_from_pdf(file):
+    loader = PyPDFLoader(file)
+    pages = loader.load_and_split()
+    text_splitter = RecursiveCharacterTextSplitter(chunk_size=250, chunk_overlap=50)
+    texts = text_splitter.split_documents(pages)
+    final_text = ""
+    for text in texts:
+        final_text = final_text + text.page_content
+    print(f"Length of final text: {len(final_text)}")
+    with open("final_text.txt", "w") as f:
+        f.write(final_text)
+    research_paper = ""
+    for line in final_text.split("\n"):
+        if line.startswith("REFERENCES"):
+            break
+        else:
+            research_paper = research_paper + line + " " 
+    with open("research_paper.txt", "w") as f:
+        f.write(research_paper)
+    print(f"Length of research paper: {len(research_paper)}")
+    return research_paper
+
+def generate_prompt(final_text):
+    prompt = f'''
+    You have been provided with a research paper in text format. Your task is to generate a mindmap structure in markdown format that summarizes the research paper.
+    Your output should use the language \"en\" 0.3 times the length of the original research paper. Do not include anything in the response, that is not the part of mindmap and use the following template (any node in the mindmap should not exceed 10-12 words, also generate additional headings that aren't present in document if required for elaborative explaination):
+        # {{Title}} (should be the title of the research paper)
+        ## {{Subtitle01}} (as required and as many as required in markdown format)
+        - {{Emoji01}} Bulletpoint01 (as required and as many as required in markdown format)
+            - {{Emoji01.1}} Bulletpoint01.1 (as required and as many as sub levels required in markdown format)
+                - {{Emoji01.1.1}} Bulletpoint01.1.1 (as required and as many as sub levels required in markdown format)
+                - {{Emoji01.1.2}} Bulletpoint01.1.2 (as required and as many as sub levels required in markdown format)
+            - {{Emoji01.2}} Bulletpoint01.2 (as required and as many as sub levels required in markdown format)
+        - {{Emoji02}} Bulletpoint02 (as required and as many as required in markdown format)
+            - {{Emoji02.1}} Bulletpoint02.1 (as required and as many as sub levels required in markdown format)
+            - {{Emoji02.2}} Bulletpoint02.2 (as required and as many as sub levels required in markdown format)
+                - {{Emoji02.2.1}} Bulletpoint02.2.1 (as required and as many as sub levels required in markdown format)
+                - {{Emoji02.2.2}} Bulletpoint02.2.2 (as required and as many as sub levels required in markdown format)
+                - {{Emoji02.2.3}} Bulletpoint02.2.3 (as required and as many as sub levels required in markdown format)
+                - {{Emoji02.2.4}} Bulletpoint02.2.4 (as required and as many as sub levels required in markdown format)
+        ## {{Subtitle02}} (as required and as many as required in markdown format)
+        - {{Emoji03}} Bulletpoint03 (as required and as many as required in markdown format)
+            - {{Emoji03.1}} Bulletpoint03.1 (as required and as many as sub levels required in markdown format)
+            - {{Emoji03.2}} Bulletpoint03.2 (as required and as many as sub levels required in markdown format)
+                - {{Emoji03.2.1}} Bulletpoint03.2.1 (as required and as many as sub levels required in markdown format)
+                - {{Emoji03.2.2}} Bulletpoint03.2.2 (as required and as many as sub levels required in markdown format)
+        - {{Emoji04}} Bulletpoint04 (as required and as many as required in markdown format)
+            - {{Emoji04.1}} Bulletpoint04.1 (as required and as many as sub levels required in markdown format)
+                - {{Emoji04.1.1}} Bulletpoint04.1.1 (as required and as many as sub levels required in markdown format)
+                - {{Emoji04.1.2}} Bulletpoint04.1.2 (as required and as many as sub levels required in markdown format)
+            - {{Emoji04.2}} Bulletpoint04.2 (as required and as many as sub levels required in markdown format)
+                - {{Emoji04.2.1}} Bulletpoint04.2.1 (as required and as many as sub levels required in markdown format)
+                - {{Emoji04.2.2}} Bulletpoint04.2.2 (as required and as many as sub levels required in markdown format)
+        Summarize the text \"{final_text}\" to generate a elaborated hierarchical mindmap structure (any node in the mindmap should not exceed 10-12 words, also generate additional headings that aren't present in document if required for elaborative explaination) markdown using the \"en\" language 0.3 times the length of the original research paper. Do not include anything in the response, that is not the part of mindmap
+    '''
+    return prompt
+
+def generate_mindmap_structure(llm, prompt):
+    response = llm.create_chat_completion(
+        messages = [
+            {'role':'user',
+            'content': prompt}
+        ],
+        temperature=0.7,
+        top_k=200,
+        top_p=3.0,
+    )
+    mindmap_data = response['choices'][0]['message']['content']
+    print(mindmap_data)
+    return mindmap_data
+
+def generate_markdown(llm, file):
+    final_text = get_text_from_pdf(file)
+    prompt = generate_prompt(final_text)
+    mindmap_markdown = generate_mindmap_structure(llm, prompt)
+    return mindmap_markdown