47 lines
1.2 KiB
Python
47 lines
1.2 KiB
Python
from langchain_community.document_loaders.generic import GenericLoader
|
|
from langchain_community.document_loaders.parsers import LanguageParser
|
|
from langchain_openai import ChatOpenAI
|
|
from langchain_core.runnables import RunnablePassthrough
|
|
from langchain_core.prompts import PromptTemplate
|
|
from langchain_core.output_parsers import StrOutputParser
|
|
from langchain_community.document_loaders import AsyncHtmlLoader
|
|
from dotenv import load_dotenv
|
|
from validators import url
|
|
from time import time
|
|
|
|
"""
|
|
Comments, TODO
|
|
"""
|
|
load_dotenv()
|
|
promt = "You are a malware reverse engineer"
|
|
def get_rag_chain():
|
|
return (
|
|
{"context": retriever | format_docs, "question": RunnablePassthrough()}
|
|
| prompt
|
|
| llm
|
|
| StrOutputParser()
|
|
)
|
|
|
|
|
|
llm = ChatOpenAI(model_name="gpt-4-turbo", temperature=0)
|
|
|
|
|
|
print(
|
|
"Welcome to HW4"
|
|
)
|
|
|
|
while True:
|
|
try:
|
|
line: str = input("llm>> ")
|
|
if line:
|
|
start_time = time()
|
|
result: str = get_rag_chain.invoke(line)
|
|
end_time = time()
|
|
elapsed_time = round(end_time - start_time, 2)
|
|
print("\n", result, "\n\nElapsed time: ", elapsed_time, " seconds")
|
|
else:
|
|
break
|
|
except Exception as e:
|
|
print(e)
|
|
break
|