Update main.py

main.py

@@ -271,7 +271,7 @@ async def __query_ml_predict(qtype: QType, content: str, header: str, token_limi
                         "messages": [
                             {"role": "user", "content": prompt.gen_prompt_wh(content=content, header=header, num_qs=num_qs, lang=l)}
                         ],
-                        "max_tokens": max(token_limit, 4096),
+                        "max_tokens": 4096,
                         "presence_penalty":0.3,
                         "temperature":0.55
                     }),
@@ -300,7 +300,7 @@ async def __query_ml_predict(qtype: QType, content: str, header: str, token_limi
                         "messages": [
                             {"role": "user", "content": prompt.gen_prompt_statements(content=content, header=header, num_qs=num_qs, lang=l)}
                         ],
-                        "max_tokens": max(token_limit, 4096),
+                        "max_tokens": 4096,
                         
                     }),
                     timeout=None
@@ -328,7 +328,7 @@ async def __query_ml_predict(qtype: QType, content: str, header: str, token_limi
                         "messages": [
                             {"role": "user", "content": prompt.gen_prompt_statements_false(content=_r_content, lang=l)}
                         ],
-                        "max_tokens": max(token_limit, 4096),
+                        "max_tokens": 4096,
                         
                     }),
                     timeout=None
@@ -532,16 +532,25 @@ async def generate_questions(request):
         if (o['style'] == QType.WH):
             
             raw_qa_list = []
-            raw_segmented: list[str] = list(filter(lambda x: (len(x)>0), o['content'].split("\n\n")))[1:]
-            print(raw_segmented)
-            for i in range(len(raw_segmented)):
-                if (len(raw_segmented[i]) and raw_segmented[i].count('\n') < 5):
-                    raw_segmented[i] += f'\n{raw_segmented[i+1]}'
-                    raw_segmented[i+1] = ""
-
-            print(raw_segmented)
-            
-            __parsed_outputs.extend(await parse_wh_question(raw_segmented, pgph_i))
+            # raw_segmented: list[str] = list(filter(lambda x: (len(x)>0), o['content'].split("\n\n")))[1:]
+
+            # for i in range(len(raw_segmented)):
+            #     if (len(raw_segmented[i]) and raw_segmented[i].count('\n') < 5):
+            #         raw_segmented[i] += f'\n{raw_segmented[i+1]}'
+            #         raw_segmented[i+1] = ""
+
+            # print(raw_segmented)
+            seg_index = 0
+            seg_index_map = ['Q.', 'A.', 'B.', 'C.', 'D.', '{']
+            raw_segmented = []
+            for seg in o['content'].split('\n'):
+                if seg.startswith(seg_index_map[seg_index]):
+                    raw_segmented.append(seg)
+                    seg_index+=1
+                if seg_index == 5:
+                    __parsed_outputs.extend(await parse_wh_question(raw_segmented, pgph_i))
+                    break
+            seg_index = 0
 
         elif (o['style'] == QType.STMT):
             print(o['content'])