import json
import requests
import os
import boto3
import shutil
from tqdm import tqdm

client = boto3.client(
    'bedrock-runtime', 
    region_name='us-west-2',
    aws_access_key_id='',
    aws_secret_access_key='',
)

model_name = 'meta.llama3-1-405b-instruct-v1:0'
malaymmlu = 'MalayMMLU_0shot.json'

folder = model_name.replace('.', '-').replace(':', '-').replace(' ', '-')
os.makedirs(folder, exist_ok = True)

if not os.path.exists(malaymmlu):
    r = requests.get('https://raw.githubusercontent.com/UMxYTL-AI-Labs/MalayMMLU/refs/heads/main/data/MalayMMLU_0shot.json')
    with open(malaymmlu, 'w') as fopen:
        fopen.write(r._content.decode())
        
with open(malaymmlu) as fopen:
    data = json.load(fopen)
    
for i in tqdm(range(len(data))):
    filename = os.path.join(folder, f'{i}.json')
    if os.path.exists(filename) and os.stat(filename).st_size > 0:
        continue
    subject = data[i]['subject']
    ques = data[i]['prompt']
    p = f"Berikut adalah soalan aneka pilihan tentang {subject}. Sila berikan jawapan sahaja.\n\n" + ques + "\nJawapan:" 
    response = client.converse(
        modelId=model_name,
        inferenceConfig= {
            'maxTokens': 10,
        },
        messages=[
            {
                'role': 'user',
                'content': [
                    {
                        'text': p,
                    }
                ]
            }
        ]
    )
    r = response['output']['message']['content'][0]['text']
    with open(filename, 'w') as fopen:
        json.dump(r, fopen)
    
shutil.make_archive(folder, 'zip', folder)