taewon99 commited on
Commit
13b9071
β€’
1 Parent(s): 01f6829

Upload app.py

Browse files
Files changed (1) hide show
  1. app.py +180 -0
app.py ADDED
@@ -0,0 +1,180 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import streamlit as st
2
+ from dotenv import load_dotenv
3
+ from langchain.text_splitter import (
4
+ CharacterTextSplitter,
5
+ RecursiveCharacterTextSplitter,
6
+ )
7
+ from langchain.vectorstores import FAISS
8
+ from langchain.embeddings import (
9
+ HuggingFaceEmbeddings,
10
+ ) # General embeddings from HuggingFace models.
11
+ from langchain.memory import ConversationBufferMemory
12
+ from langchain.chains import ConversationalRetrievalChain
13
+ from htmlTemplates import css, bot_template, user_template
14
+ from langchain.llms import LlamaCpp # For loading transformer models.
15
+ from langchain.document_loaders import PyPDFLoader, TextLoader, JSONLoader, CSVLoader
16
+ import tempfile # μž„μ‹œ νŒŒμΌμ„ μƒμ„±ν•˜κΈ° μœ„ν•œ λΌμ΄λΈŒλŸ¬λ¦¬μž…λ‹ˆλ‹€.
17
+ import os
18
+ from huggingface_hub import hf_hub_download # Hugging Face Hubμ—μ„œ λͺ¨λΈμ„ λ‹€μš΄λ‘œλ“œν•˜κΈ° μœ„ν•œ ν•¨μˆ˜μž…λ‹ˆλ‹€.
19
+
20
+
21
+ # PDF λ¬Έμ„œλ‘œλΆ€ν„° ν…μŠ€νŠΈλ₯Ό μΆ”μΆœν•˜λŠ” ν•¨μˆ˜μž…λ‹ˆλ‹€.
22
+ def get_pdf_text(pdf_docs):
23
+ temp_dir = tempfile.TemporaryDirectory() # μž„μ‹œ 디렉토리λ₯Ό μƒμ„±ν•©λ‹ˆλ‹€.
24
+ temp_filepath = os.path.join(temp_dir.name, pdf_docs.name) # μž„μ‹œ 파일 경둜λ₯Ό μƒμ„±ν•©λ‹ˆλ‹€.
25
+ with open(temp_filepath, "wb") as f: # μž„μ‹œ νŒŒμΌμ„ λ°”μ΄λ„ˆλ¦¬ μ“°κΈ° λͺ¨λ“œλ‘œ μ—½λ‹ˆλ‹€.
26
+ f.write(pdf_docs.getvalue()) # PDF λ¬Έμ„œμ˜ λ‚΄μš©μ„ μž„μ‹œ νŒŒμΌμ— μ”λ‹ˆλ‹€.
27
+ pdf_loader = PyPDFLoader(temp_filepath) # PyPDFLoaderλ₯Ό μ‚¬μš©ν•΄ PDFλ₯Ό λ‘œλ“œν•©λ‹ˆλ‹€.
28
+ pdf_doc = pdf_loader.load() # ν…μŠ€νŠΈλ₯Ό μΆ”μΆœν•©λ‹ˆλ‹€.
29
+ return pdf_doc # μΆ”μΆœν•œ ν…μŠ€νŠΈλ₯Ό λ°˜ν™˜ν•©λ‹ˆλ‹€.
30
+
31
+
32
+ # 과제
33
+ # μ•„λž˜ ν…μŠ€νŠΈ μΆ”μΆœ ν•¨μˆ˜λ₯Ό μž‘μ„±
34
+ def get_text_file(docs):
35
+ temp_dir = tempfile.TemporaryDirectory() # μž„μ‹œ 디렉토리λ₯Ό μƒμ„±ν•©λ‹ˆλ‹€.
36
+ temp_filepath = os.path.join(temp_dir.name, docs.name) # μž„μ‹œ 파일 경둜λ₯Ό μƒμ„±ν•©λ‹ˆλ‹€.
37
+ with open(temp_filepath, "wb") as f: # μž„μ‹œ νŒŒμΌμ„ λ°”μ΄λ„ˆλ¦¬ μ“°κΈ° λͺ¨λ“œλ‘œ μ—½λ‹ˆλ‹€.
38
+ f.write(docs.getvalue()) # ν…μŠ€νŠΈ 파일의 λ‚΄μš©μ„ μž„μ‹œ νŒŒμΌμ— μ”λ‹ˆλ‹€.
39
+ text_loader = TextLoader(temp_filepath) # TextLoaderλ₯Ό μ‚¬μš©ν•΄ ν…μŠ€νŠΈλ₯Ό λ‘œλ“œν•©λ‹ˆλ‹€.
40
+ text_doc = text_loader.load() # ν…μŠ€νŠΈλ₯Ό μΆ”μΆœν•©λ‹ˆλ‹€.
41
+ return text_doc # μΆ”μΆœν•œ ν…μŠ€νŠΈλ₯Ό λ°˜ν™˜ν•©λ‹ˆλ‹€.
42
+
43
+
44
+ def get_csv_file(docs):
45
+ temp_dir = tempfile.TemporaryDirectory() # μž„μ‹œ 디렉토리λ₯Ό μƒμ„±ν•©λ‹ˆλ‹€.
46
+ temp_filepath = os.path.join(temp_dir.name, docs.name) # μž„μ‹œ 파일 경둜λ₯Ό μƒμ„±ν•©λ‹ˆλ‹€.
47
+ with open(temp_filepath, "wb") as f: # μž„μ‹œ νŒŒμΌμ„ λ°”μ΄λ„ˆλ¦¬ μ“°κΈ° λͺ¨λ“œλ‘œ μ—½λ‹ˆλ‹€.
48
+ f.write(docs.getvalue()) # CSV 파일의 λ‚΄μš©μ„ μž„μ‹œ νŒŒμΌμ— μ”λ‹ˆλ‹€.
49
+ csv_loader = CSVLoader(temp_filepath) # CSVLoaderλ₯Ό μ‚¬μš©ν•΄ CSVλ₯Ό λ‘œλ“œν•©λ‹ˆλ‹€.
50
+ csv_doc = csv_loader.load() # ν…μŠ€νŠΈλ₯Ό μΆ”μΆœν•©λ‹ˆλ‹€.
51
+ return csv_doc # μΆ”μΆœν•œ ν…μŠ€νŠΈλ₯Ό λ°˜ν™˜ν•©λ‹ˆλ‹€.
52
+
53
+
54
+ def get_json_file(docs):
55
+ temp_dir = tempfile.TemporaryDirectory() # μž„μ‹œ 디렉토리λ₯Ό μƒμ„±ν•©λ‹ˆλ‹€.
56
+ temp_filepath = os.path.join(temp_dir.name, docs.name) # μž„μ‹œ 파일 경둜λ₯Ό μƒμ„±ν•©λ‹ˆλ‹€.
57
+ with open(temp_filepath, "wb") as f: # μž„μ‹œ νŒŒμΌμ„ λ°”μ΄λ„ˆλ¦¬ μ“°κΈ° λͺ¨λ“œλ‘œ μ—½λ‹ˆλ‹€.
58
+ f.write(docs.getvalue()) # JSON 파일의 λ‚΄μš©μ„ μž„μ‹œ νŒŒμΌμ— μ”λ‹ˆλ‹€.
59
+ json_loader = JSONLoader(temp_filepath) # JSONLoaderλ₯Ό μ‚¬μš©ν•΄ JSON을 λ‘œλ“œν•©λ‹ˆλ‹€.
60
+ json_doc = json_loader.load() # ν…μŠ€νŠΈλ₯Ό μΆ”μΆœν•©λ‹ˆλ‹€.
61
+ return json_doc # μΆ”μΆœν•œ ν…μŠ€νŠΈλ₯Ό λ°˜ν™˜ν•©λ‹ˆλ‹€.
62
+
63
+
64
+ # λ¬Έμ„œλ“€μ„ μ²˜λ¦¬ν•˜μ—¬ ν…μŠ€νŠΈ 청크둜 λ‚˜λˆ„λŠ” ν•¨μˆ˜μž…λ‹ˆλ‹€.
65
+ def get_text_chunks(documents):
66
+ text_splitter = RecursiveCharacterTextSplitter(
67
+ chunk_size=1000, # 청크의 크기λ₯Ό μ§€μ •ν•©λ‹ˆλ‹€.
68
+ chunk_overlap=200, # 청크 μ‚¬μ΄μ˜ 쀑볡을 μ§€μ •ν•©λ‹ˆλ‹€.
69
+ length_function=len, # ν…μŠ€νŠΈμ˜ 길이λ₯Ό μΈ‘μ •ν•˜λŠ” ν•¨μˆ˜λ₯Ό μ§€μ •ν•©λ‹ˆλ‹€.
70
+ )
71
+
72
+ documents = text_splitter.split_documents(documents) # λ¬Έμ„œλ“€μ„ 청크둜 λ‚˜λˆ•λ‹ˆλ‹€.
73
+ return documents # λ‚˜λˆˆ 청크λ₯Ό λ°˜ν™˜ν•©λ‹ˆλ‹€.
74
+
75
+
76
+ # ν…μŠ€νŠΈ μ²­ν¬λ“€λ‘œλΆ€ν„° 벑터 μŠ€ν† μ–΄λ₯Ό μƒμ„±ν•˜λŠ” ν•¨μˆ˜μž…λ‹ˆλ‹€.
77
+ def get_vectorstore(text_chunks):
78
+ # μ›ν•˜λŠ” μž„λ² λ”© λͺ¨λΈμ„ λ‘œλ“œν•©λ‹ˆλ‹€.
79
+ embeddings = HuggingFaceEmbeddings(
80
+ model_name="bert-base-multilingual-cased",
81
+ model_kwargs={"device": "cpu"},
82
+ ) # μž„λ² λ”© λͺ¨λΈμ„ μ„€μ •ν•©λ‹ˆλ‹€.
83
+ vectorstore = FAISS.from_documents(text_chunks, embeddings) # FAISS 벑터 μŠ€ν† μ–΄λ₯Ό μƒμ„±ν•©λ‹ˆλ‹€.
84
+ return vectorstore # μƒμ„±λœ 벑터 μŠ€ν† μ–΄λ₯Ό λ°˜ν™˜ν•©λ‹ˆλ‹€.
85
+
86
+
87
+ def get_conversation_chain(vectorstore):
88
+ model_name_or_path = "TheBloke/Llama-2-7B-chat-GGUF"
89
+ model_basename = "llama-2-7b-chat.Q2_K.gguf"
90
+ model_path = hf_hub_download(repo_id=model_name_or_path, filename=model_basename)
91
+
92
+ llm = LlamaCpp(
93
+ model_path=model_path,
94
+ n_ctx=4086,
95
+ input={"temperature": 0.75, "max_length": 2000, "top_p": 1},
96
+ verbose=True,
97
+ )
98
+ # λŒ€ν™” 기둝을 μ €μž₯ν•˜κΈ° μœ„ν•œ λ©”λͺ¨λ¦¬λ₯Ό μƒμ„±ν•©λ‹ˆλ‹€.
99
+ memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
100
+ # λŒ€ν™” 검색 체인을 μƒμ„±ν•©λ‹ˆλ‹€.
101
+ conversation_chain = ConversationalRetrievalChain.from_llm(
102
+ llm=llm, retriever=vectorstore.as_retriever(), memory=memory
103
+ )
104
+ return conversation_chain # μƒμ„±λœ λŒ€ν™” 체인을 λ°˜ν™˜ν•©λ‹ˆλ‹€.
105
+
106
+
107
+ # μ‚¬μš©μž μž…λ ₯을 μ²˜λ¦¬ν•˜λŠ” ν•¨μˆ˜μž…λ‹ˆλ‹€.
108
+ def handle_userinput(user_question):
109
+ print("user_question => ", user_question)
110
+ # λŒ€ν™” 체인을 μ‚¬μš©ν•˜μ—¬ μ‚¬μš©μž μ§ˆλ¬Έμ— λŒ€ν•œ 응닡을 μƒμ„±ν•©λ‹ˆλ‹€.
111
+ response = st.session_state.conversation({"question": user_question})
112
+ # λŒ€ν™” 기둝을 μ €μž₯ν•©λ‹ˆλ‹€.
113
+ st.session_state.chat_history = response["chat_history"]
114
+
115
+ for i, message in enumerate(st.session_state.chat_history):
116
+ if i % 2 == 0:
117
+ st.write(
118
+ user_template.replace("{{MSG}}", message.content),
119
+ unsafe_allow_html=True,
120
+ )
121
+ else:
122
+ st.write(
123
+ bot_template.replace("{{MSG}}", message.content), unsafe_allow_html=True
124
+ )
125
+
126
+
127
+ def main():
128
+ load_dotenv()
129
+ # Streamlit UIλ₯Ό μ„€μ •ν•©λ‹ˆλ‹€.
130
+ st.title("λ¬Έμ„œμ— λŒ€ν•΄ μ§ˆλ¬Έν•˜λŠ” μ±„νŒ…")
131
+ st.sidebar.title("λ¬Έμ„œ μ—…λ‘œλ“œ")
132
+ st.sidebar.text("여기에 μ—¬λŸ¬λΆ„μ˜ λ¬Έμ„œλ₯Ό μ—…λ‘œλ“œν•˜μ„Έμš”.")
133
+ st.sidebar.text("PDF, TXT, CSV, JSON 파일 ν˜•μ‹μ„ μ§€μ›ν•©λ‹ˆλ‹€.")
134
+ st.write(css, unsafe_allow_html=True)
135
+
136
+ if "conversation" not in st.session_state:
137
+ st.session_state.conversation = None
138
+ if "chat_history" not in st.session_state:
139
+ st.session_state.chat_history = None
140
+
141
+ st.header("λ¬Έμ„œμ— λŒ€ν•΄ μ§ˆλ¬Έν•˜λŠ” μ±„νŒ…")
142
+ user_question = st.text_input("문선에 λŒ€ν•΄μ„œ μ§ˆλ¬Έμ„ ν•΄λ³΄μ„Έμš”")
143
+ if user_question:
144
+ handle_userinput(user_question)
145
+
146
+ with st.sidebar:
147
+ st.subheader("Your documents")
148
+ docs = st.file_uploader("여기에 λ¬Έμ„œλ₯Ό μ—…λ‘œλ“œν•˜κ³  '처리'λ₯Ό ν΄λ¦­ν•˜μ„Έμš”", accept_multiple_files=True)
149
+ if st.button("처리"):
150
+ with st.spinner("Processing"):
151
+ # get pdf text
152
+ doc_list = []
153
+
154
+ for file in docs:
155
+ print("file - type : ", file.type)
156
+ if file.type == "text/plain":
157
+ # file is .txt
158
+ doc_list.extend(get_text_file(file))
159
+ elif file.type in ["application/octet-stream", "application/pdf"]:
160
+ # file is .pdf
161
+ doc_list.extend(get_pdf_text(file))
162
+ elif file.type == "text/csv":
163
+ # file is .csv
164
+ doc_list.extend(get_csv_file(file))
165
+ elif file.type == "application/json":
166
+ # file is .json
167
+ doc_list.extend(get_json_file(file))
168
+
169
+ # get the text chunks
170
+ text_chunks = get_text_chunks(doc_list)
171
+
172
+ # create vector store
173
+ vectorstore = get_vectorstore(text_chunks)
174
+
175
+ # create conversation chain
176
+ st.session_state.conversation = get_conversation_chain(vectorstore)
177
+
178
+
179
+ if __name__ == "__main__":
180
+ main()