LLM / llm_generator.py
Phaneendra99's picture
Update llm_generator.py
de89c02 verified
from langchain import PromptTemplate
from langchain.prompts.chat import (
ChatPromptTemplate,
SystemMessagePromptTemplate,
AIMessagePromptTemplate,
HumanMessagePromptTemplate,
)
import os
from langchain.memory import ConversationBufferMemory
from langchain.chains import ConversationChain
from langchain.chains import ConversationalRetrievalChain
from langchain.memory import ChatMessageHistory,ConversationSummaryBufferMemory
# LLM Generator
from langchain.llms import CTransformers
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
# Get the directory path of the current script
script_dir = os.path.dirname(os.path.abspath(__file__))
config = {'max_new_tokens': 256,
'temperature': 0.4,
'repetition_penalty': 1.1,
'context_length': 4096, # Set to max for Chat Summary, Llama-2 has a max context length of 4096
}
llm = CTransformers(model=os.path.join(script_dir, 'transformers', 'llama-2-7b-chat.Q2_K.gguf'),
callbacks=[StreamingStdOutCallbackHandler()],
config=config)
# Define system and user message templates
with open(os.path.join(script_dir, 'prompts', 'system_message_template.txt'), 'r') as file:
system_message_template = file.read().replace('\n', '')
with open(os.path.join(script_dir, 'prompts', 'user_mesage_template.txt'), 'r') as file:
user_message_template = file.read().replace('\n', '')
with open(os.path.join(script_dir, 'prompts', 'condense_question_prompt.txt'), 'r') as file:
condense_question_prompt = file.read().replace('\n', '')
# Create message templates
system_message = SystemMessagePromptTemplate.from_template(system_message_template)
user_message = HumanMessagePromptTemplate.from_template(user_message_template)
# Compile messages into a chat prompt template
messages = [system_message, user_message]
chatbot_prompt = ChatPromptTemplate.from_messages(messages)
# array_of_files
from rag_pipeline import instantiate_rag
retriever = instantiate_rag()
history = ChatMessageHistory()
# Provide the chat history when initializing the ConversationalRetrievalChain
qa = ConversationalRetrievalChain.from_llm(
llm,
retriever=retriever,
memory = ConversationSummaryBufferMemory(
memory_key="chat_history",
input_key="question",
llm=llm,
max_token_limit=40,
return_messages=True
),
return_source_documents=False,
chain_type="stuff",
max_tokens_limit=100,
condense_question_prompt= PromptTemplate.from_template(condense_question_prompt),
combine_docs_chain_kwargs={'prompt': chatbot_prompt},
verbose=True,
return_generated_question=False,
)
def LLM_generator(question: str):
answer = qa({"question": question,"chat_history":history.messages})["answer"]
print("##------##")
return answer
# Implement Classification
from nlp_models import sentiment_class, pattern_classification, corelation_analysis
# is_depressed = sentiment_class(conversation_buffer)
# random
# Initialize the 2D list
is_depressed = [[]]
# Assign a probability value to the cell at row 0 and column 1
is_depressed[0].append(0.75)
# Check if the probability value is greater than 0.5 and print the result
if is_depressed[0][0] > 0.5:
print('Not so depressed')
else:
print('is_depressed')