gchat_processing/article_test_script.py - master - poly.tex/pAPer

gchat_processing/article_test_script.py

54 строки2 KB

Danila

master: semscholar & gigachat clients refactored with async http requests

28 апр 2024, 20:56

28 апр 2024, 20:56c8c3993

import asyncio

from gch_chain_wrapper import GigaChainWrapper
from pathlib import Path
import os
import json


async def main():
    with open(Path('credentials/gigachat_config.json')) as fp:
        giga_token = json.load(fp=fp)['credentials']
    gchain_wrapper = GigaChainWrapper(
        giga_token=giga_token,
        sql_cache_path=Path('sqlite_cache.db'),
        giga_scope='GIGACHAT_API_CORP'
    )
    with open('test_data/source_texts/test_idea.txt', 'r', encoding='utf-8') as file:
        idea_text = file.read()
    print('Idea:')
    print(idea_text)
    print('Queries:')
    queries_list = await gchain_wrapper.get_search_queries(idea_text, n_queries=3)
    print(queries_list)
    print('Articles:')
    art_files = os.listdir('test_data/articles')
    summaries = []
    for filename in art_files:
        with open(f'test_data/articles/{filename}', 'r', encoding='utf-8') as file:
            d = json.load(file)
        print('Current abstract:')
        print(d['annotation'])
        print('Summarized:')
        summary = await gchain_wrapper.get_summary_simple(d['annotation'])
        print(summary)
        summaries.append((d['title'], summary))
    print('Generated annotation:')
    giga_annot = gchain_wrapper.get_annotation(idea_text, summaries)
    print(giga_annot)
    print('Generated themes:')
    giga_themes = await gchain_wrapper.get_themes(idea_text, summaries, 4)
    for theme in giga_themes:
        print('......')
        print(theme)
        print('......')
    print('Generated related works:')
    giga_rel_works = await gchain_wrapper.get_rel_works(idea_text, summaries, giga_themes)
    for work in giga_rel_works:
        print('......')
        print(work)
        print('......')


if __name__ == '__main__':
    asyncio.run(main())