"""Pilot: wyszukiwanie po polsku w korpusie 90 tekstów serwisu (1633 fragmenty). Przebieg 2 (7.10.2026 wieczorem):
poprawione wejście dokumentów. W przebiegu 1 akapity EmbeddingGemma 2 dostawały prompt_name="Document" (w konfiguracji modelu
"title: none | text: ") ORAZ ręczny format "title: {tytuł} | text: {akapit}", czyli podwójny przedrostek, przed którym karta modelu
ostrzega. Teraz dokument to sam tekst "title: {tytuł} | text: {akapit}" bez prompt_name; pytania nadal z prompt_name="SearchQuery".
Dla równych warunków e5 także dostaje tytuł: "passage: {tytuł}. {akapit}" (w przebiegu 1 sam akapit). BM25 bez zmian.
Miary zamrożone przed uruchomieniem: Hit@5, Recall@5, MRR@10 na poziomie dokumentu
(fragment -> dokument przez max podobieństwa), osobno dla pytań FAQ (łatwe, 423)
i pytań potocznych (trudne, 60). Systemy: EmbeddingGemma 2 (768d i 256d), multilingual-e5-small, BM25."""
import json, time, os, sys, resource, hashlib, platform, re
import numpy as np
# Rewizje modeli przypięte 8.10.2026; przebiegi 2 i 3 uruchomiono na tych samych rewizjach (pobranych 7.10.2026) bez tego argumentu.
REV={'intfloat/multilingual-e5-small':'614241f622f53c4eeff9890bdc4f31cfecc418b3','google/embeddinggemma-2':'914f7f89142e33e77833254d9c9b90c3cef7303b'}
P=os.path.dirname(os.path.abspath(__file__))
K=P+'/korpus'
docs=json.load(open(K+'/fragmenty.json',encoding='utf-8'))
faq=json.load(open(K+'/pytania-faq.json',encoding='utf-8'))
hard=json.load(open(K+'/pytania-trudne.json',encoding='utf-8'))
frag_doc=[d['doc'] for d in docs]
def sha(p): return hashlib.sha256(open(p,'rb').read()).hexdigest()[:16]
def rss_mb(): return resource.getrusage(resource.RUSAGE_SELF).ru_maxrss/1048576
def metrics(scores, rel_sets):
    """scores: (n_q, n_frag); rel_sets: list of sets of doc ids"""
    out=[]
    for i,rel in enumerate(rel_sets):
        order=np.argsort(-scores[i]); seen=[]; 
        for j in order:
            d=frag_doc[j]
            if d not in seen: seen.append(d)
            if len(seen)>=10: break
        top5=seen[:5]
        hit=any(d in rel for d in top5)
        rec=sum(1 for d in rel if d in top5)/len(rel)
        rr=0.0
        for r,d in enumerate(seen,1):
            if d in rel: rr=1.0/r; break
        out.append({'hit5':hit,'recall5':rec,'mrr10':rr,'top5':top5})
    return out
def summarize(rows):
    n=len(rows); return {'n':n,'hit@5':round(sum(r['hit5'] for r in rows)/n,4),'recall@5':round(sum(r['recall5'] for r in rows)/n,4),'mrr@10':round(sum(r['mrr10'] for r in rows)/n,4)}
qsets={'faq':([q['q'] for q in faq],[{q['doc']} for q in faq]),'trudne':([q['q'] for q in hard],[set(q['doc']) for q in hard])}
results={'data':time.strftime('%Y-%m-%d %H:%M'),'host':platform.platform(),'python':sys.version.split()[0],
 'korpus':{'fragmenty':len(docs),'dokumenty':len(set(frag_doc)),'sha':{f:sha(K+'/'+f) for f in ['fragmenty.json','pytania-faq.json','pytania-trudne.json']}},'systemy':{}}
perq=open(P+'/wyniki-per-pytanie.jsonl','w',encoding='utf-8')
def record(name, scores, extra):
    res={}
    for sname,(qs,rels) in qsets.items():
        rows=metrics(scores[sname],rels); res[sname]=summarize(rows)
        for q,rel,r in zip(qs,rels,rows): perq.write(json.dumps({'system':name,'zestaw':sname,'q':q,'rel':sorted(rel),**r},ensure_ascii=False)+'\n')
    results['systemy'][name]={**extra,**res}; print(name,json.dumps(res,ensure_ascii=False),flush=True)
    json.dump(results,open(P+'/wyniki.json','w',encoding='utf-8'),ensure_ascii=False,indent=1)
# --- BM25
from rank_bm25 import BM25Okapi
tok=lambda s: re.findall(r'\w+',s.lower())
t0=time.time(); bm=BM25Okapi([tok(d['title']+' '+d['text']) for d in docs]); t_idx=time.time()-t0
sc={}; lat=[]
for sname,(qs,_) in qsets.items():
    arr=[]
    for q in qs:
        t=time.time(); arr.append(bm.get_scores(tok(q))); lat.append(time.time()-t)
    sc[sname]=np.array(arr)
record('bm25',sc,{'indeks_s':round(t_idx,2),'zapytanie_ms_mediana':round(1000*float(np.median(lat)),1),'rss_mb':round(rss_mb())})
# --- embeddingi
import torch
from sentence_transformers import SentenceTransformer
device='mps' if torch.backends.mps.is_available() else 'cpu'
def run_embed(name, model_id, q_prompt, d_prompt, dims=None):
    t0=time.time(); m=SentenceTransformer(model_id, device=device, revision=REV.get(model_id)); t_load=time.time()-t0
    prompts=getattr(m,'prompts',{}) or {}
    def enc(texts, kind, bs=32):
        pn = q_prompt if kind=='q' else d_prompt
        if pn is None: return m.encode(texts, batch_size=bs, normalize_embeddings=True, convert_to_numpy=True).astype(np.float32)
        if pn in prompts: return m.encode(texts, prompt_name=pn, batch_size=bs, normalize_embeddings=True, convert_to_numpy=True).astype(np.float32)
        return m.encode([pn+t for t in texts], batch_size=bs, normalize_embeddings=True, convert_to_numpy=True).astype(np.float32)
    dtexts=[f"title: {d['title']} | text: {d['text']}" if 'gemma' in model_id else f"{d['title']}. {d['text']}" for d in docs]
    t0=time.time(); D=enc(dtexts,'d'); t_idx=time.time()-t0; assert np.isfinite(D).all()
    for dim in dims or [None]:
        Dd=D if dim is None else D[:,:dim]/np.linalg.norm(D[:,:dim],axis=1,keepdims=True)
        sc={}; lat=[]
        for sname,(qs,_) in qsets.items():
            t=time.time(); Q=enc(qs,'q'); lat.append((time.time()-t)/len(qs))
            Qd=Q if dim is None else Q[:,:dim]/np.linalg.norm(Q[:,:dim],axis=1,keepdims=True)
            sc[sname]=Qd@Dd.T
        label=name if dim is None else f'{name}-{dim}d'
        record(label,sc,{'model':model_id,'device':device,'wymiary':int(Dd.shape[1]),'prompty':{'q':q_prompt,'q_tekst':prompts.get(q_prompt,q_prompt),'d':d_prompt,'format_dokumentu':('title: {tytul} | text: {akapit}' if 'gemma' in model_id else 'passage: {tytul}. {akapit}'),'z_modelu':q_prompt in prompts},'przebieg':2,'ladowanie_s':round(t_load,1),'indeks_s':round(t_idx,1),'fragmenty_na_s':round(len(docs)/t_idx,1),'zapytanie_ms_srednia':round(1000*float(np.mean(lat)),1),'rss_mb':round(rss_mb()),'sentence_transformers':__import__('sentence_transformers').__version__,'torch':torch.__version__})
    del m
run_embed('e5-small','intfloat/multilingual-e5-small','query: ','passage: ')
run_embed('embeddinggemma-2','google/embeddinggemma-2','SearchQuery',None,dims=[None,256])
perq.close(); print('KONIEC')