from datasets import load_dataset
from langchain.chains import LLMChain
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate

from langchain import hub
from langchain.retrievers.multi_query import MultiQueryRetriever
from langchain.docstore.document import Document
import os



# os.environ["DASHSCOPE_API_KEY"] = "sk-cc1c8314fdbd43ceaf26ec1824d5dd3b"
# llm = Tongyi()
# data = load_dataset("jamescalam/ai-arxiv-chunked", split="train")
# data = data[:2]
# print(data)


numbers = [[1,2], [3,4], [4,5]]
squared_numbers = map(lambda x: x[0]**2, numbers)
print(list(squared_numbers))



# docs = []
# for row in data:
#     doc = Document(
#         page_content=row["chunk"],
#         metadata={
#             "title": row["title"],
#             "source": row["source"],
#             "id": row["id"],
#             "chunk-id": row["chunk-id"],
#             "text": row["chunk"]
#         }
#     )
#     docs.append(doc)


#     embeddings = JinaEmbeddings(
#     jina_api_key="jina_7e2c88997a50417aab497c15a4c6cec7vuBoG_CK-_0gYILG38ZIoJHTL1_q", model_name="jina-embeddings-v2-base-en"
# )
    

# vectorstore = Chroma.from_documents(docs, embeddings)

# retriever = vectorstore.as_retriever()

# retriever = MultiQueryRetriever.from_llm(
#     retriever= retriever, llm=llm
# )

    # print(docs)