app.py

app.py

@@ -131,8 +131,8 @@ print(f"Input file path: {INPUT_FILE}")
 print(f"Output file path: {OUTPUT_FILE}")
 PROMPT_ID = args.prompt_id if args.prompt_id else INPUT_FILE
 PROMPT_TOPIC = args.prompt_topic if args.prompt_topic else INPUT_FILE
-MULTI_PV = args.multi_pv if args.multi_pv else 1000
-NUM_WORDS = args.num_words if args.num_words else 10000
+MULTI_PV = args.multi_pv if args.multi_pv else 100
+NUM_WORDS = args.num_words if args.num_words else 20
 NUM_TOKENS = args.num_tokens if args.num_tokens else 10000
 BEAM_WIDTH = args.beam_width if args.beam_width else 1
 ALPHA_MODE = args.alpha_mode if args.alpha_mode else 0
@@ -232,6 +232,7 @@ def main(INPUT_FILE, OUTPUT_FILE, PROMPT_ID, PROMPT_TOPIC, MULTI_PV, NUM_WORDS,
             f.write("The quick brown fox jumps over the lazy dog")
 
     # OUTPUT_FILE = "output.lif"  # Changed output filename
+    textoutput = ""
     MODEL_CONTEXT_WINDOW = 128_000  # Example context window, adjust if needed for the actual model
     SAFETY_THRESHOLD = 2_000
     MAX_INPUT_TOKENS = MODEL_CONTEXT_WINDOW - SAFETY_THRESHOLD  # Max tokens per model *input slice*
@@ -360,6 +361,7 @@ def main(INPUT_FILE, OUTPUT_FILE, PROMPT_ID, PROMPT_TOPIC, MULTI_PV, NUM_WORDS,
         with open(OUTPUT_FILE, 'w', encoding='utf-8') as writer:
             print("  Writing header to output file...")
             writer.write(''.join(header_lines))
+            textoutput += ''.join(header_lines)
             print("  Header written. Starting word-by-word prediction.\n")
 
             for turnCount in range(START_TURN, START_TURN + num_words_to_process):
@@ -526,6 +528,7 @@ def main(INPUT_FILE, OUTPUT_FILE, PROMPT_ID, PROMPT_TOPIC, MULTI_PV, NUM_WORDS,
                 # 9.7 Write to file
                 # print("  9.7: Writing lines to output file...")
                 writer.write(''.join(lines))
+                textoutput += ''.join(lines)
                 # print("    Write complete.\n")
 
                 # 9.8 Update state
@@ -581,6 +584,8 @@ def main(INPUT_FILE, OUTPUT_FILE, PROMPT_ID, PROMPT_TOPIC, MULTI_PV, NUM_WORDS,
         torch.cuda.empty_cache()
     print("\nScript finished.")
 
+    return textoutput
+
 
 ### RUN MAIN ####
 
@@ -588,7 +593,8 @@ def main(INPUT_FILE, OUTPUT_FILE, PROMPT_ID, PROMPT_TOPIC, MULTI_PV, NUM_WORDS,
      #START_TURN, MODEL_TAG)
 
 def premain(ip):
-   main("Kangaroos.txt", "KangaroosByQ06.lif", "KangarooEssay", "The Mating Habits of Kangaroos", 100, 1000, 10000, 1, 0, 1, "Qwen")
+    txtoutput = main("Kangaroos.txt", "KangaroosByQ06.lif", "KangarooEssay", "The Mating Habits of Kangaroos", 100, 1000, 10000, 1, 0, 1, "Qwen")
+    return txtoutput
 
 demo = gr.Interface(fn=premain, inputs="text", outputs="text")
 #demo = gr.Interface(fn=premain)