Team Ai
Apppublic

Multimedika/Bot_Development

sourceHugging Facemitupdated 2y agoView on Hugging Face
0likes
get_topic.py83 linesDownload Raw Back to script
1import nest_asyncio2import os3from dotenv import load_dotenv4from jinja2 import Template5from pydantic import BaseModel, Field6from pymongo.mongo_client import MongoClient7 8from llama_index.program.openai import OpenAIPydanticProgram9from llama_index.core.extractors import PydanticProgramExtractor10from llama_index.llms.openai import OpenAI11 12from core.prompt import ADD_METADATA_TEMPLATE13from core.summarization.summarizer import SummarizeGenerator14 15nest_asyncio.apply()16 17load_dotenv()18 19 20class NodeMetadata(BaseModel):21    """Metadata for nodes, capturing topic and subtopic from the book."""22 23    topic: str = Field(24        ...,25        description="The main subject or category that the node is associated with, representing a broad theme within the book.",26    )27    subtopic: str = Field(28        ...,29        description="A more specific aspect or section under the main topic, refining the context of the node within the book.",30    )31 32 33def extract_topic(references, content_table):34    uri = os.getenv("MONGO_URI")35    client = MongoClient(uri)36    37    try:38        client.admin.command('ping')39        print("Pinged your deployment. You successfully connected to MongoDB!")40    except Exception as e:41        print(e)42        # Access a specific database43    db = client["summarizer"]44 45    # Access a collection within the database46    collection = db["topic_collection"]47        48    generate_content_table = SummarizeGenerator(references)49    extractor_output, extractor_dics  = generate_content_table.extract_content_table(content_table)50    print(extractor_output)51    data_to_insert = {52    "title": references["title"],53    **extractor_dics  # Unpack the extractor_output dictionary54    }55    56    collection.insert_one(data_to_insert)57    58 59    add_metadata_template = str(60        Template(ADD_METADATA_TEMPLATE).render(extractor_output=extractor_output)61    )62 63    print("add metadata template : ", add_metadata_template)64 65    llm = OpenAI(temperature=0.1, model="gpt-4o-mini")66 67    openai_program = OpenAIPydanticProgram.from_defaults(68        output_cls=NodeMetadata,69        prompt_template_str="{input}",70        extract_template_str=add_metadata_template,71        llm=llm,72    )73 74    topic_extractor = PydanticProgramExtractor(75        program=openai_program,76        input_key="input",77        show_progress=True,78        extract_template_str=add_metadata_template,79        llm=llm,80    )81 82    return topic_extractor83