from datasets import load_dataset
from openai import OpenAI
ds = load_dataset("RTVS/SpotifyLyrics001")


client = OpenAI(
  base_url="https://openrouter.ai/api/v1",
  api_key="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",
)
summarized_lyrics = []

for index in range(len(ds['train']['text'])):
    message = ds['train']['text'][index]
    response = client.chat.completions.create(
    extra_body={},
    model="openrouter/optimus-alpha",
    messages=[
        {
        "role": "user",
        "content": [
            {
            "type": "text",
            "text": f"summarize this lyrics of the song: {message}"
            },
        ]
        }
    ]
    ).choices[0].message.content
    print("ok")
    summarized_lyrics.append(response)
ds = ds.add_column("summarized-lyrics", summarized_lyrics)
ds.push_to_hub("tqhuyen/Spotify-Lyrics-Summarized")