llm/llm.py

# {'prompts': [
#    {'id': 1, 'prompt_template': 'Question: {question}\nAnswer: \n', 'description': 'simple question without a prompt', 'rate': 1}, 
#    {'id': 2, 'prompt_template': "Question: {question}\nAnswer: Write a concise answer on the question with 
#               one example if it's possible. CONCISE ANSWER.\n", 'description': 'simple concise prompt', 'rate': 3}]}


import os

import yaml
import logging

from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain 
from langchain.llms import HuggingFaceHub
from .config import config 

class LLM_chain:
    def __init__(self):
        self.llm = HuggingFaceHub(
            repo_id=config["model"], 
            model_kwargs={"temperature": config["temperature"], "max_new_tokens": config["max_new_tokens"], "top_k": config["top_k"], "load_in_8bit": config["load_in_8bit"]})

    def __read_yaml(self):
        try:
            # get current dir 
            current_dir = os.path.dirname(os.path.realpath(__file__))
            yaml_file = os.path.join(current_dir, 'prompts.yaml')
            
            with open(yaml_file, 'r') as file:
                data = yaml.safe_load(file)
            return data
        except Exception as e:
            logging.error(e)

    def __call__(self, entity: str, id: int = 0):
        try:
            data = self.__read_yaml()
            prompts = data["prompts"][id] # get second prompt from yaml, need change id parameter to get other prompt
            template = prompts["prompt_template"]
            prompt = PromptTemplate(template=template, input_variables=["entity"])
            llm_chain = LLMChain(prompt=prompt, llm=self.llm, verbose=True)
            output = llm_chain.invoke(entity)
            return output["text"]
        except Exception as e:
            logging.error(e)