plnlvv/tag-magic

1import os2import re3import json4import time5from pathlib import Path6from threading import Thread7 8import pandas as pd9import streamlit as st10import torch11from transformers import AutoTokenizer, AutoModelForCausalLM12 13BASE_MODEL = os.getenv("BASE_MODEL", "Qwen/Qwen2.5-1.5B-Instruct")14MAX_NEW_TOKENS = int(os.getenv("MAX_NEW_TOKENS", "512"))15QUANT_MODE = os.getenv("QUANT_MODE", "none")16 17st.set_page_config(page_title="Interview Coding Assistant", layout="wide")18st.title("Автоматическая разметка интервью")19st.markdown("---")20 21if "generating" not in st.session_state:22    st.session_state.generating = False23if "result" not in st.session_state:24    st.session_state.result = None25if "error" not in st.session_state:26    st.session_state.error = None27 28def build_prompt_h2(example, tokenizer):29    instruction = f"""You are an expert in interview analysis. Your task is to highlight the thematic codes in the interview text and corresponding quotes.30 31Act step by step:32 331. Carefully review the output format shown in the example below. Remember that each code block must contain "Общий код" (General code), then Quote, then "Конкретный код" (Specific code). The quote must be verbatim and enclosed in quotation marks.342. Read the interview transcript and the topic. Identify all fragments (quotes) that relate to the interview topic.353. Group the quotes by general themes — these will be the "Общий код" (General codes). For each general theme, come up with a short name.364. Within each general code, identify specific meaning aspects — these will be the "конкретный код" (Specific codes). The names of specific codes should reflect the essence of the quote.375. Generate the answer strictly following the format from the example. Do not add any explanations, do not write words like 'Step 1', 'Step 2' — only the final blocks of codes and quotes.38 39Format of an output (consists of several general codes, each followed by quotes and specific codes):40**Общий код 1: <generate general code 1>**41"<quote text>" - **<generate specific code 1> (Конкретный код)**42"<quote text>" - **<generate specific code 2> (Конкретный код)**43**Общий код 2: <generate general code 2>**44"<quote text>" - **<generate specific code> (Конкретный код)**45**Общий код <general code number>: <generate general code>**46"<quote text>" - **<generate specific code> (Конкретный код)**47and so on. You choose the number of general and specific codes.48 49Now you should do the markup for the interview according to the plan. Important: The answer should contain only codes and quotes, without unnecessary words and repetitions.50The quotes must be strictly from the text. Give the answer in Russian.51"""52 53    user_content = f"""Тема интервью:54{example['topic']}55 56Текст интервью:57{example['transcript']}58 59Пожалуйста, выполни разметку интервью, каждый общий код в указанном формате:60**Общий код <generate general code>: <general code name>**61"<quote text>" - **<generate specific code> (Конкретный код)**"""62 63    messages = [64        {"role": "system", "content": instruction},65        {"role": "user", "content": user_content},66    ]67    prompt = tokenizer.apply_chat_template(68        messages,69        tokenize=False,70        add_generation_prompt=True,71    )72    return prompt73 74@st.cache_resource(show_spinner="Загрузка модели...")75def load_model():76    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True)77    if tokenizer.pad_token is None:78        tokenizer.pad_token = tokenizer.eos_token79    if tokenizer.pad_token_id is None:80        tokenizer.pad_token_id = tokenizer.eos_token_id81    82    model = AutoModelForCausalLM.from_pretrained(83        BASE_MODEL,84        device_map="auto",85        trust_remote_code=True,86        torch_dtype=torch.float16,87        low_cpu_mem_usage=True88    )89    model.eval()90    return tokenizer, model91 92def generate(93    tokenizer,94    model,95    topic: str,96    transcript: str,97    max_new_tokens: int,98    repetition_penalty: float,99    no_repeat_ngram_size: int100) -> str:101    example = {102        "topic": topic.strip(),103        "transcript": transcript.strip()104    }105    106    prompt = build_prompt_h2(example, tokenizer)107    108    inputs = tokenizer(109        prompt,110        return_tensors="pt",111        truncation=True,112        max_length=2048113    ).to(model.device)114    115    with torch.inference_mode():116        outputs = model.generate(117            **inputs,118            max_new_tokens=max_new_tokens,119            do_sample=False,120            num_beams=1,121            repetition_penalty=repetition_penalty,122            no_repeat_ngram_size=no_repeat_ngram_size,123            temperature=1.0,124            top_p=1.0,125            top_k=0,126            early_stopping=False,127            pad_token_id=tokenizer.pad_token_id,128            eos_token_id=tokenizer.eos_token_id,129        )130    131    generated = tokenizer.decode(132        outputs[0][inputs["input_ids"].shape[1]:],133        skip_special_tokens=True134    ).strip()135    136    for token in ["<|im_end|>", "<|im_start|>", "<|endoftext|>"]:137        generated = generated.split(token)[0].strip()138    139    return generated140 141def parse_output(raw_text: str) -> list:142    rows = []143    current_general = ""144    lines = raw_text.split("\n")145    146    for line in lines:147        line = line.strip()148        if not line:149            continue150        151        general_pattern = r'\*{0,2}Общий код\s*\d+\s*[:\-]\s*(.+?)\*{0,2}$'152        match_general = re.match(general_pattern, line, re.IGNORECASE)153        if match_general:154            current_general = match_general.group(1).strip().strip('*').strip()155            continue156        157        quote_pattern = r'^["“«](.+?)["”»]\s*[-–—]\s*\*{0,2}(.+?)\*{0,2}$'158        match_quote = re.match(quote_pattern, line)159        if match_quote and current_general:160            rows.append({161                "Общий код": current_general,162                "Конкретный код": match_quote.group(2).strip(),163                "Цитата": match_quote.group(1).strip(),164            })165        elif current_general and (line.startswith('"') or line.startswith('“') or line.startswith('«')):166            quote = line.strip('"“”«»').strip()167            if quote:168                rows.append({169                    "Общий код": current_general,170                    "Конкретный код": "",171                    "Цитата": quote,172                })173    return rows174 175def validate_parsed_data(rows: list) -> bool:176    if not rows:177        return False178    for row in rows:179        if not row.get("Общий код") or not row.get("Цитата"):180            return False181    return True182 183with st.sidebar:184    st.header("Параметры генерации")185    186    repetition_penalty = st.slider(187        "repetition_penalty",188        min_value=1.0,189        max_value=2.0,190        value=1.5,191        step=0.1,192        help="Штраф за повторения"193    )194    195    no_repeat_ngram_size = st.slider(196        "no_repeat_ngram_size",197        min_value=0,198        max_value=8,199        value=4,200        help="Размер n-граммы, которую нельзя повторять"201    )202    203    max_new_tokens = st.slider(204        "max_new_tokens",205        min_value=128,206        max_value=512,207        value=256,208        step=64,209        help="Максимальное количество генерируемых токенов"210    )211    212    st.divider()213    st.subheader("Информация о модели")214    st.caption(f"**Модель:** {BASE_MODEL.split('/')[-1]}")215    st.caption(f"**Квантизация:** {QUANT_MODE}")216    st.divider()217    st.markdown("""218    **Инструкция:** 219    1. Введите тему исследования220    2. Вставьте транскрипт интервью (до 3000 символов)221    3. Нажмите "Разметить"222    4. Подождите 1-2 минуты223    5. Скачайте результат224    """)225 226topic = st.text_area(227    "Тема исследования",228    height=100,229    placeholder="Например: Поколенческая идентичность, жизненные выборы и устойчивость российской молодежи",230)231 232transcript = st.text_area(233    "Текст интервью",234    height=300,235    placeholder="Вставьте транскрипт интервью сюда...\n\nИнтервьюер: Вопрос...\nИнформант: Ответ...",236)237 238if len(transcript) > 3000:239    st.warning(f"Текст сокращён с {len(transcript)} до 3000 символов")240    transcript = transcript[:3000]241 242col1, col2, col3 = st.columns([1, 1, 4])243run_btn = col1.button("Разметить", type="primary", use_container_width=True, 244                      disabled=st.session_state.generating)245clear_btn = col2.button("Очистить", use_container_width=True)246 247if clear_btn:248    st.session_state.result = None249    st.session_state.error = None250    st.rerun()251 252try:253    tokenizer, model = load_model()254except Exception as e:255    st.error(f"Ошибка загрузки модели: {str(e)}")256    st.stop()257 258if run_btn and not st.session_state.generating:259    if not topic.strip():260        st.error("Укажите тему исследования.")261    elif not transcript.strip():262        st.error("Вставьте текст интервью.")263    else:264        st.session_state.generating = True265        st.session_state.result = None266        st.session_state.error = None267        st.rerun()268 269if st.session_state.generating:270    status_placeholder = st.empty()271    status_placeholder.info("Генерация разметки... Это может занять 1-2 минуты. Пожалуйста, не обновляйте страницу.")272    273    progress_bar = st.progress(0)274    progress_bar.progress(25)275    276    try:277        progress_bar.progress(50)278        raw_output = generate(279            tokenizer, model, topic, transcript,280            max_new_tokens, repetition_penalty, no_repeat_ngram_size281        )282        283        progress_bar.progress(75)284        parsed_rows = parse_output(raw_output)285        286        if parsed_rows and validate_parsed_data(parsed_rows):287            st.session_state.result = {288                "raw": raw_output,289                "parsed": parsed_rows290            }291            progress_bar.progress(100)292            status_placeholder.success("Генерация завершена!")293        else:294            st.session_state.error = "Не удалось распарсить структуру ответа"295            status_placeholder.error(st.session_state.error)296            297    except Exception as e:298        st.session_state.error = str(e)299        status_placeholder.error(f"Ошибка: {str(e)}")300    301    st.session_state.generating = False302    st.rerun()303 304# Отображение результатов305if st.session_state.result:306    st.subheader("Структурированная разметка")307    308    df = pd.DataFrame(st.session_state.result["parsed"])309    310    st.dataframe(311        df,312        use_container_width=True,313        column_config={314            "Общий код": st.column_config.TextColumn("Общий код", width="medium"),315            "Конкретный код": st.column_config.TextColumn("Конкретный код", width="large"),316            "Цитата": st.column_config.TextColumn("Цитата", width="large"),317        }318    )319    320    st.info(f"Размечено: {len(df)} фрагментов | {df['Общий код'].nunique()} уникальных общих кодов")321    322    col_dl1, col_dl2 = st.columns(2)323    with col_dl1:324        st.download_button(325            "Скачать CSV",326            df.to_csv(index=False).encode("utf-8-sig"),327            "interview_coding.csv",328            "text/csv",329        )330    with col_dl2:331        st.download_button(332            "Скачать JSON",333            json.dumps(st.session_state.result["parsed"], ensure_ascii=False, indent=2).encode("utf-8"),334            "interview_coding.json",335            "application/json",336        )337    338    with st.expander("Сырой ответ модели"):339        st.code(st.session_state.result["raw"], language="markdown")340    341    if st.button("Новая разметка"):342        st.session_state.result = None343        st.rerun()344 345if st.session_state.error and not st.session_state.result:346    st.error(f"Ошибка: {st.session_state.error}")347    if st.button("Попробовать снова"):348        st.session_state.error = None349        st.rerun()350 351st.markdown("---")352st.caption("Совет: Для лучших результатов используйте транскрипты объёмом до 3000 символов")