ice_breaking_challenge/background_task.py

import time
import random
import re

from flask import session
import pandas as pd

from ice_breaking_challenge import model, quiz_generated, quiz_results
from ice_breaking_challenge.google_sheets import QUESTIONS

def background_task(sid, qna: pd.DataFrame, team_size: int):
    """Session에 저장된 응답을 바탕으로 Gemma 모델을 사용한 퀴즈 생성"""
    import time
    start_time = time.time()

    global quiz_results
    quiz_generated[sid] = False
    generated_quizzes = generate_quiz(qna, team_size)
    quiz_results[sid] = generated_quizzes
    # 퀴즈 생성 완료
    quiz_generated[sid] = True

    end_time = time.time()
    print(end_time - start_time)


def sample_quizzes(data: pd.DataFrame, team_size: int):
    questions_size = min(10 // team_size + 1, 5)
    sample_indices: dict[str, list[str]] = dict()
    for row in data.itertuples():
        sample_indices[row.Index] = random.sample(['Q1', 'Q2', 'Q3', 'Q4', 'Q5'], questions_size)

    sample_quizzes = []
    for key, value in sample_indices.items():
        for question_number in value:
            question = QUESTIONS.iloc[0][question_number]
            answer = data.iloc[key][question_number]
            name = data.iloc[key]['name']
            sample_quizzes.append([question, answer, name])

    return sample_quizzes[:2]  # TODO


def generate_quiz(data: pd.DataFrame, team_size: int):
    global model
    generated_quizzes = []

    # selected_data = random.sample(data, min(10, len(data)))
    selected_data = sample_quizzes(data, team_size)

    template_input="""
        <instruction>
        Using the text: {question} {answer}, create a new multiple-choice question with 4 answer options.
        """
    for row in selected_data:
        question, answer, name = row

        input_text = template_input.format(question=question, answer=answer)

        response = model.generate(input_text, max_length=512)
        print(response)
        lines = response.split('\n')
        lines = [line.strip() for line in lines]

        opening_tag_counter = 0
        line_index_to_remove = None
        for i, line in enumerate(lines):
            if bool(re.search('^<[^/]', line)):
                opening_tag_counter += 1

            if opening_tag_counter == 2:
                line_index_to_remove = i
                break

        lines = lines[line_index_to_remove + 1:]
        lines = [re.sub('<.\w+>', "", line) for line in lines]
        lines = [line.strip() for line in lines]
        lines = [line for line in lines if line]

        print(lines)
        print("==========================================")

        question_generated = lines[0]
        multiple_choice_generated = lines[1]
        answer_generated = lines[2]

        print(question_generated, multiple_choice_generated, answer_generated)
        print("==========================================")
        generated_quizzes.append([name, question_generated, multiple_choice_generated, answer_generated])

    return generated_quizzes