Update app.py

app.py

@@ -1,8 +1,5 @@
 #!/usr/bin/env python3
 
-
-
-
 import os, json, time, random, threading, logging
 from datetime import datetime, timezone
 import torch; torch.set_num_threads(os.cpu_count()); torch.set_num_interop_threads(os.cpu_count())
@@ -18,77 +15,139 @@ TOKENS_PER_PROMPT = 2048
 SECS_PER_TOKEN = 15
 TEMP = 0.9; TOP_P = 0.95; MAX_CTX = 8192
 
-
 logging.basicConfig(level=logging.INFO)
 log = logging.getLogger()
 
-def _rj(p,d):
-    
-    try: return json.load(open(p,encoding="utf-8"))
-    except: return d
-        
+def _rj(p, d):
+    try:
+        return json.load(open(p, encoding="utf-8"))
+    except:
+        return d
 
-def _aw(p,o):
-    t=p+".tmp"; open(t,"w",encoding="utf-8").write(json.dumps(o,ensure_ascii=False,indent=2)); os.replace(t,p)
 
-prompts=_rj(PROMPTS_PATH,[])
-if not prompts: raise Exception("no prompts")
+def _aw(p, o):
+    t = p + ".tmp"
+    open(t, "w", encoding="utf-8").write(json.dumps(o, ensure_ascii=False, indent=2))
+    os.replace(t, p)
 
-tok=os.environ.get("HF_READ_TOKEN")
-log.info("loading model...")
-tokenizer=AutoTokenizer.from_pretrained(MODEL_NAME,token=tok)
-model=AutoModelForCausalLM.from_pretrained(MODEL_NAME,torch_dtype=torch.float32,low_cpu_mem_usage=False,token=tok)
-model.to("cpu");model.eval()
-log.info("model up")
+prompts = _rj(PROMPTS_PATH, [])
+if not prompts:
+    raise Exception("No prompts found in full_prompts.json")
 
-lock=threading.Lock()
+tok = os.environ.get("HF_READ_TOKEN")
+log.info("Loading model...")
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, token=tok)
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_NAME,
+    torch_dtype=torch.float32,
+    low_cpu_mem_usage=False,
+    token=tok
+)
+model.to("cpu"); model.eval()
+log.info("Model is ready.")
+
+lock = threading.Lock()
 
 def _init():
-    s=_rj(STATE_PATH,{})
-    if not s or s.get("finished"):
-        i=random.randrange(len(prompts))
-        s={"i":i,"p":prompts[i],"g":"","c":0,"t":time.time(),"finished":False}
-        _aw(STATE_PATH,s)
-    return s
-
-def _es(st):
-    d=int(time.time()-st);h,r=divmod(d,3600);m,s=divmod(r,60)
+    state = _rj(STATE_PATH, {})
+    if not state or state.get("finished"):
+        idx = random.randrange(len(prompts))
+        state = {"i": idx, "p": prompts[idx], "g": "", "c": 0, "t": time.time(), "finished": False}
+        _aw(STATE_PATH, state)
+    return state
+
+def _es(start_time):
+    elapsed = int(time.time() - start_time)
+    h, rem = divmod(elapsed, 3600)
+    m, s = divmod(rem, 60)
     return f"{h}h {m}m {s}s"
 
 def _loop():
     while True:
-        with lock: s=_init()
-        if s["finished"]: time.sleep(SECS_PER_TOKEN); continue
-        c=s["p"]+s["g"]
-        ids=tokenizer(c,return_tensors="pt",truncation=True,max_length=MAX_CTX).input_ids
-        with torch.no_grad(): out=model.generate(ids,max_new_tokens=1,do_sample=True,temperature=TEMP,top_p=TOP_P)
-        nt=tokenizer.decode(out[0,-1],skip_special_tokens=True,clean_up_tokenization_spaces=False)
         with lock:
-            s["g"]+=nt; s["c"]+=1
-            if s["c"]>=TOKENS_PER_PROMPT: s["finished"]=True
-            _aw(STATE_PATH,s)
+            st = _init()
+        if st["finished"]:
+            time.sleep(SECS_PER_TOKEN)
+            continue
+        context = st["p"] + st["g"]
+        ids = tokenizer(context, return_tensors="pt", truncation=True, max_length=MAX_CTX).input_ids
+        with torch.no_grad():
+            out = model.generate(
+                ids,
+                max_new_tokens=1,
+                do_sample=True,
+                temperature=TEMP,
+                top_p=TOP_P
+            )
+        next_token = tokenizer.decode(out[0, -1], skip_special_tokens=True, clean_up_tokenization_spaces=False)
+        with lock:
+            st["g"] += next_token
+            st["c"] += 1
+            if st["c"] >= TOKENS_PER_PROMPT:
+                st["finished"] = True
+            _aw(STATE_PATH, st)
         time.sleep(SECS_PER_TOKEN)
-threading.Thread(target=_loop,daemon=True).start()
+
+threading.Thread(target=_loop, daemon=True).start()
 
 def _fetch():
-    s=_rj(STATE_PATH,{})
-    if not s: return "...","","0h 0m 0s"
-    return s["p"],s["g"],_es(s["t"])
-
-def _sg(f,i):
-    f1,i1=f.strip(),i.strip()
-    if not(f1 or i1): return gr.update(value="eh?"),gr.update(),gr.update()
-    p,g,e=_fetch();guess=f1 or i1;gt="full" if f1 else "idea"
-    r={"ts":datetime.now(timezone.utc).isoformat(),"prompt":p,"time":e,"resp":g,"guess":guess,"type":gt}
-    with lock: open(DATA_PATH,"a",encoding="utf-8").write(json.dumps(r,ensure_ascii=False)+"\n")
-    return gr.update(value="ok logged"),gr.update(value=""),gr.update(value="")
+    state = _rj(STATE_PATH, {})
+    if not state:
+        return "...", "", "0h 0m 0s"
+    return state["p"], state["g"], _es(state["t"])
+
+def _submit_prediction(detailed, summary):
+    det = detailed.strip()
+    if not det:
+        return gr.update(value="Please enter at least a detailed prediction."), gr.update(value=""), gr.update(value="")
+    prompt_text, oracle_resp, elapsed = _fetch()
+    record = {
+        "ts": datetime.now(timezone.utc).isoformat(),
+        "prompt": prompt_text,
+        "time": elapsed,
+        "resp": oracle_resp,
+        "prediction": det,
+        "summary": summary.strip()
+    }
+    with lock:
+        open(DATA_PATH, "a", encoding="utf-8").write(json.dumps(record, ensure_ascii=False) + "\n")
+    return gr.update(value="Prediction logged!"), gr.update(value=""), gr.update(value="")
 
 with gr.Blocks(theme="darkdefault") as demo:
-    gr.Markdown("# What Comes Next")
-    prm=gr.Markdown();txt=gr.Textbox(lines=10,interactive=False,label="oracle");tme=gr.Textbox(interactive=False,label="time")
-    rbtn=gr.Button("refresh");full=gr.Textbox(label="full");idea=gr.Textbox(label="idea");send=gr.Button("send");st=gr.Textbox(interactive=False,label="status")
-    demo.load(_fetch,outputs=[prm,txt,tme])
-    rbtn.click(_fetch,outputs=[prm,txt,tme])
-    send.click(_sg,inputs=[full,idea],outputs=[st,full,idea])
-
-if __name__=="__main__": demo.launch(server_name="0.0.0.0",server_port=7860)
+    gr.Markdown(
+        "# What Comes Next\n"
+        "Enter what you think will come next in the text.\n"
+        "Provide a detailed continuation and optionally a brief summary for context."
+    )
+    prompt_md = gr.Markdown()
+    oracle_output = gr.Textbox(lines=10, interactive=False, label="Oracle Response")
+    time_info = gr.Textbox(interactive=False, label="Elapsed Time")
+
+    with gr.Row():
+        prompt_md, oracle_output, time_info
+
+    detailed = gr.Textbox(
+        label="Your Detailed Prediction",
+        placeholder="Enter the full text continuation you expect...",
+        lines=3
+    )
+    summary = gr.Textbox(
+        label="Prediction Summary (Optional)",
+        placeholder="Optionally, summarize your prediction in a few words...",
+        lines=2
+    )
+    status = gr.Textbox(interactive=False, label="Status")
+    submit_btn = gr.Button("Submit Prediction")
+    refresh_btn = gr.Button("Refresh Oracle")
+
+    demo.load(_fetch, outputs=[prompt_md, oracle_output, time_info])
+    refresh_btn.click(_fetch, outputs=[prompt_md, oracle_output, time_info])
+    submit_btn.click(
+        _submit_prediction,
+        inputs=[detailed, summary],
+        outputs=[status, detailed, summary]
+    )
+
+if __name__ == "__main__":
+    demo.launch(server_name="0.0.0.0", server_port=7860)
+