darknet_diaries_llm/main.py
Romain Quinet 3d734a3064 Chat mode
2023-10-06 23:22:10 +02:00

70 lines
2.4 KiB
Python

from llama_index import (SimpleDirectoryReader, ServiceContext, StorageContext, PromptTemplate,
load_index_from_storage, Document, set_global_service_context)
from llama_index.node_parser import SimpleNodeParser
from llama_index import VectorStoreIndex
from llama_index.llms import OpenAI, ChatMessage, MessageRole
from llama_index.prompts import PromptTemplate
from llama_index.chat_engine.condense_question import CondenseQuestionChatEngine
import os
import re
llm = OpenAI(model="gpt-4", temperature=0, max_tokens=256)
service_context = ServiceContext.from_defaults(llm=llm)
set_global_service_context(service_context)
if not os.path.exists("./index/lock"):
print("Generating index...")
documents = []
for filename in os.listdir("./data"):
episode_number = re.search(r'\d+', filename).group()
with open("./data/" + filename, 'r') as f:
content = f.read()
document = Document(
text=content,
metadata={
"episode_number": episode_number
}
)
documents = SimpleDirectoryReader('./data').load_data()
parser = SimpleNodeParser.from_defaults()
nodes = parser.get_nodes_from_documents(documents)
index = VectorStoreIndex(nodes, show_progress=True)
index.storage_context.persist(persist_dir="./index")
open("./index/lock", 'a').close()
else:
print("Loading index...")
storage_context = StorageContext.from_defaults(persist_dir="./index")
index = load_index_from_storage(storage_context)
custom_prompt = PromptTemplate(
"You have been trained on the Darknet Diaries podcast transcripts with data from october 6 2023."
"You are now an expert about it and will answer as such. You know about every episode up to number 138. \n"
"----------------\n"
"Chat history: {chat_history}\n"
"----------------\n"
"Please answer this question by referring to the podcast: {question}"
)
custom_chat_history = []
query_engine = index.as_query_engine()
chat_engine = CondenseQuestionChatEngine.from_defaults(
query_engine=query_engine,
condense_question_prompt=custom_prompt,
chat_history=custom_chat_history,
verbose=True
)
while True:
try:
user_prompt = input("Prompt: ")
streaming_response = chat_engine.stream_chat(user_prompt)
for token in streaming_response.response_gen:
print(token, end="")
print("\n")
except KeyboardInterrupt:
break