156 lines
4.7 KiB
Python
156 lines
4.7 KiB
Python
import argparse
|
|
import json
|
|
import os
|
|
import sys
|
|
|
|
import pandas as pd
|
|
import streamlit as st
|
|
|
|
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
|
|
|
|
from dataset.search_shortqa_dataset import normalize_answers, normalize_contexts, normalize_gold_doc_ids
|
|
from trainer.search_shortqa_reward import parse_search_response, score_search_shortqa_response
|
|
|
|
|
|
DEFAULT_DATA = "dataset/examples/search_shortqa_raw_demo.jsonl"
|
|
|
|
|
|
def load_jsonl(path):
|
|
rows = []
|
|
if not os.path.exists(path):
|
|
return rows
|
|
with open(path, "r", encoding="utf-8") as f:
|
|
for line in f:
|
|
line = line.strip()
|
|
if line:
|
|
rows.append(json.loads(line))
|
|
return rows
|
|
|
|
|
|
def context_table(contexts):
|
|
return pd.DataFrame(
|
|
[
|
|
{
|
|
"id": doc.get("id", ""),
|
|
"title": doc.get("title", ""),
|
|
"text": doc.get("text", ""),
|
|
}
|
|
for doc in contexts
|
|
]
|
|
)
|
|
|
|
|
|
def row_from_editor(sample, table):
|
|
row = dict(sample)
|
|
row["contexts"] = [
|
|
{
|
|
"id": str(item.get("id", "")).strip(),
|
|
"title": str(item.get("title", "")).strip(),
|
|
"text": str(item.get("text", "")).strip(),
|
|
}
|
|
for item in table.to_dict("records")
|
|
if str(item.get("id", "")).strip() and str(item.get("text", "")).strip()
|
|
]
|
|
return row
|
|
|
|
|
|
def default_prediction(sample):
|
|
answers = normalize_answers(sample)
|
|
gold_doc_ids = normalize_gold_doc_ids(sample)
|
|
return json.dumps(
|
|
{
|
|
"answer": answers[0] if answers else "不确定",
|
|
"citations": gold_doc_ids[:3],
|
|
},
|
|
ensure_ascii=False,
|
|
)
|
|
|
|
|
|
def parse_args():
|
|
parser = argparse.ArgumentParser(description="SearchShortQA Streamlit demo")
|
|
parser.add_argument("--data", default=DEFAULT_DATA)
|
|
return parser.parse_args()
|
|
|
|
|
|
def main():
|
|
args = parse_args()
|
|
st.set_page_config(page_title="SearchShortQA", layout="wide")
|
|
|
|
st.title("SearchShortQA")
|
|
|
|
with st.sidebar:
|
|
data_path = st.text_input("数据文件", value=args.data)
|
|
rows = load_jsonl(data_path)
|
|
if not rows:
|
|
st.error("数据为空")
|
|
return
|
|
sample_ids = [str(row.get("id") or i) for i, row in enumerate(rows)]
|
|
selected_id = st.selectbox("样本", sample_ids, index=0)
|
|
|
|
sample = rows[sample_ids.index(selected_id)]
|
|
sample = dict(sample)
|
|
question = st.text_input("问题", value=str(sample.get("question") or sample.get("query") or sample.get("prompt") or ""))
|
|
sample["question"] = question
|
|
|
|
left, right = st.columns([1.2, 1.0])
|
|
with left:
|
|
st.subheader("检索片段")
|
|
edited_contexts = st.data_editor(
|
|
context_table(normalize_contexts(sample)),
|
|
num_rows="dynamic",
|
|
use_container_width=True,
|
|
column_config={
|
|
"id": st.column_config.TextColumn("ID", width="small"),
|
|
"title": st.column_config.TextColumn("标题", width="medium"),
|
|
"text": st.column_config.TextColumn("正文", width="large"),
|
|
},
|
|
)
|
|
sample = row_from_editor(sample, edited_contexts)
|
|
|
|
with right:
|
|
st.subheader("模型输出")
|
|
prediction = st.text_area("JSON", value=default_prediction(sample), height=180)
|
|
reward, parts = score_search_shortqa_response(
|
|
prediction,
|
|
answers=normalize_answers(sample),
|
|
gold_doc_ids=normalize_gold_doc_ids(sample),
|
|
contexts=normalize_contexts(sample),
|
|
)
|
|
answer, citations, is_json = parse_search_response(prediction)
|
|
|
|
c1, c2, c3, c4 = st.columns(4)
|
|
c1.metric("Reward", f"{reward:.3f}")
|
|
c2.metric("F1", f"{parts['answer_f1']:.3f}")
|
|
c3.metric("引用合法率", f"{parts['citation_valid']:.3f}")
|
|
c4.metric("格式", f"{parts['format_score']:.3f}")
|
|
|
|
st.subheader("奖励分解")
|
|
st.dataframe(
|
|
pd.DataFrame([parts]).T.rename(columns={0: "score"}),
|
|
use_container_width=True,
|
|
)
|
|
|
|
st.subheader("引用检查")
|
|
valid_ids = {doc["id"] for doc in normalize_contexts(sample)}
|
|
cited = set(citations)
|
|
citation_rows = []
|
|
for doc in normalize_contexts(sample):
|
|
citation_rows.append(
|
|
{
|
|
"id": doc["id"],
|
|
"title": doc["title"],
|
|
"cited": doc["id"] in cited,
|
|
"gold": doc["id"] in set(normalize_gold_doc_ids(sample)),
|
|
"valid": doc["id"] in valid_ids,
|
|
"text": doc["text"],
|
|
}
|
|
)
|
|
st.dataframe(pd.DataFrame(citation_rows), use_container_width=True)
|
|
|
|
st.subheader("答案")
|
|
st.write(answer if is_json else prediction)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|