def generate_chat_responses()

in genai-on-vertex-ai/gemini/model_upgrades/multiturn_chat/vertex_script/eval.py [0:0]


def generate_chat_responses(project_id: str, location:str, model: str, dataset: pd.DataFrame, response_column_name: str) -> None:
    '''Generate the final model response for each conversation in the dataset using the specified model.'''
    client = genai.Client(vertexai=True, project=project_id, location=location)
    responses = []
    user_prompts = []
    for i, record in dataset.iterrows():
        print(f'Generating chat completion #{i+1} with {model}')
        messages = json.loads(record.get('history'))
        last_user_message = messages.pop()
        history = [
            Content(
                role=message['role'],
                parts=[Part(text=message['content'])],
            )
            for message in messages
        ]
        chat = client.chats.create(model=model, history=history)
        response = chat.send_message(message=[Part(text=last_user_message['content'])])
        user_prompts.append(last_user_message)
        responses.append( response.candidates[0].content.parts[0].text )
    dataset['prompt'] = user_prompts  # The last user message is required by the Autorater
    dataset[response_column_name] = responses 
    print(f'{len(responses)} responses from model {model} are stored in dataset column "{response_column_name}"')