1
0
mirror of https://github.com/microsoft/qlib.git synced 2026-07-01 01:51:18 +08:00
Files
qlib/tests/finco/test_sumarize.py
Fivele-Li 1d88830b0d Add recorder task and visualize (#1542)
* add recorder task

* add batch generate summarize report unittest.

* * add recorder to RecorderTask;
* add matplot figure to analyzer.py

* add image to markdown;

* Add some log

* update figure path.

---------

Co-authored-by: Young <afe.young@gmail.com>
Co-authored-by: Cadenza-Li <362237642@qq.com>
2023-06-12 15:48:00 +08:00

67 lines
2.0 KiB
Python

import unittest
import os
import shutil
from dotenv import load_dotenv
# pydantic support load_dotenv, so load_dotenv will be deprecated in the future.
from qlib.finco.task import SummarizeTask
from qlib.finco.workflow import WorkflowContextManager
from qlib.finco.llm import APIBackend
from qlib.finco.workflow import WorkflowManager
load_dotenv(verbose=True, override=True)
class TestSummarize(unittest.TestCase):
def test_chat(self):
messages = [
{
"role": "system",
"content": "Your are a professional financial assistant.",
},
{
"role": "user",
"content": "How to write a perfect quant strategy.",
},
]
response = APIBackend().try_create_chat_completion(messages=messages)
print(response)
def test_execution(self):
task = SummarizeTask()
context = WorkflowContextManager()
context.set_context("workspace", "../../examples/benchmarks/Linear")
context.set_context("user_prompt", "My main focus is on the performance of the strategy's return."
"Please summarize the information and give me some advice.")
task.assign_context_manager(context)
resp = task.execute()
print(resp)
def test_generate_batch_result(self):
wm = WorkflowManager()
prompt = wm.default_user_prompt
# prompt = ""
workdir = os.path.dirname(wm.get_context().get_context("workspace"))
summaries_path = os.path.join(workdir, "summaries")
if not os.path.exists(summaries_path):
os.makedirs(summaries_path)
for i in range(10):
wm.run(prompt)
if os.path.exists(f"{workdir}/finCoReport.md"):
shutil.move(f"{workdir}/finCoReport.md", f"{workdir}/summaries/finCoReport{i}.md")
def test_parse2txt(self):
task = SummarizeTask()
resp = task.get_info_from_file("")
print(resp)
if __name__ == "__main__":
unittest.main()