jdmorzan commited on
Commit
5a01bde
·
verified ·
1 Parent(s): d0a1605

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +173 -0
app.py ADDED
@@ -0,0 +1,173 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import threading
2
+ import queue
3
+ import time
4
+
5
+ from langchain.chains import ConversationalRetrievalChain
6
+ from langchain.memory import ConversationBufferMemory
7
+ from langchain.chat_models import ChatOpenAI
8
+ from langsmith import traceable
9
+ from langchain.embeddings.openai import OpenAIEmbeddings
10
+ from langchain.vectorstores import Chroma
11
+ from langchain.prompts import ChatPromptTemplate
12
+ from langchain.prompts.chat import SystemMessagePromptTemplate, HumanMessagePromptTemplate
13
+ from langchain.callbacks.base import BaseCallbackHandler
14
+ import gradio as gr
15
+
16
+ # --------------------------
17
+ # Custom Streaming Callback Handler that uses a queue.
18
+ # --------------------------
19
+ class CustomStreamingCallbackHandler(BaseCallbackHandler):
20
+ def __init__(self):
21
+ self.token_queue = queue.Queue()
22
+
23
+ def on_llm_new_token(self, token: str, **kwargs):
24
+ # Simply put the new token in the queue.
25
+ self.token_queue.put(token)
26
+
27
+ # Instantiate the custom streaming callback (do not pass the chatbot here).
28
+ stream_handler = CustomStreamingCallbackHandler()
29
+
30
+ # --------------------------
31
+ # Setup vector database and embeddings
32
+ # --------------------------
33
+ embeddings = OpenAIEmbeddings()
34
+ vectordb = Chroma(
35
+ persist_directory="./ai_internacionalidad_v2",
36
+ embedding_function=embeddings
37
+ )
38
+
39
+ # --------------------------
40
+ # Define prompt templates
41
+ # --------------------------
42
+ general_system_template = r"""
43
+ Cuando indique qué carrera quiere estudiar responde que el 100% de las carreras de UPC cuenta con un programa internacional: https://www.youtube.com/watch?v=-TgOaphgwYg
44
+ e indica que ellos pueden vivir experiencias como las de alguno de nuestros estudiantes: https://www.youtube.com/watch?v=qFRN6ruRrfc
45
+
46
+ Luego indica que si quiere estudiar en Inglaterra, Australia, Canadá, Nueva Zelandia, Estados Unidos puede hacerlo con el programa NCUK y muestrasles el video: https://youtu.be/C9cKwDt8CwA?si=1NfbsBzaItDIJd5m
47
+
48
+ Luego, si quieren estudiar alguna carrera de las siguientes factultades: arquitectura o Comunicaciones o Derecho o Economia o Ingeniería o Negocios o Psicología, indica que si quieren estudiar desde el Perú en una universidad de Estados Unidos o en Estados Unidos también lo pueden hacer con el programa Arizona: https://youtu.be/jbvMRNEuZUA?si=6FJifohlnoge4VJb
49
+
50
+ Además, en base a la carrera indicada presenta los siguiente videos indicando conoce más sobre como se vive la internacionalización en la carrera de tu interés en caso tengas el link específico de la carrera o en la facultad en caso no tengas el link hacia la carrera:
51
+ Toma los siguientes documentos de contexto {context} y responde únicamente basado en este contexto.
52
+ """
53
+
54
+ general_user_template = "Pregunta:```{question}```"
55
+ messages = [
56
+ SystemMessagePromptTemplate.from_template(general_system_template),
57
+ HumanMessagePromptTemplate.from_template(general_user_template)
58
+ ]
59
+ qa_prompt = ChatPromptTemplate.from_messages(messages)
60
+
61
+ # --------------------------
62
+ # Create conversation memory
63
+ # --------------------------
64
+ def create_memory():
65
+ return ConversationBufferMemory(memory_key='chat_history', return_messages=True)
66
+
67
+ # --------------------------
68
+ # Define the chain function that uses the LLM to answer queries
69
+ # --------------------------
70
+ def pdf_qa(query, memory, llm):
71
+ chain = ConversationalRetrievalChain.from_llm(
72
+ llm=llm,
73
+ retriever=vectordb.as_retriever(search_kwargs={'k': 28}),
74
+ combine_docs_chain_kwargs={'prompt': qa_prompt},
75
+ memory=memory
76
+ )
77
+ return chain({"question": query})
78
+
79
+ # --------------------------
80
+ # Build the Gradio Interface with custom CSS for the "Enviar" button.
81
+ # --------------------------
82
+ with gr.Blocks() as demo:
83
+ # Inject custom CSS via HTML.
84
+ gr.HTML(
85
+ """
86
+ <style>
87
+ /* Target the button inside the container with id "enviar_button" */
88
+ #enviar_button button {
89
+ background-color: #E50A17 !important;
90
+ color: white !important;
91
+ }
92
+ </style>
93
+ """
94
+ )
95
+
96
+ # Chatbot component with an initial greeting.
97
+ chatbot = gr.Chatbot(
98
+ label="Internacionalidad",
99
+ value=[[None,
100
+ '''¡Hola!
101
+
102
+ Dime la carrera que te interesa y te contaré qué experiencia puedes vivir en el extranjero y como otros alumnos UPC ya estan viviendo esa experiencia.
103
+
104
+ ¡Hazme cualquier pregunta y descubramos juntas todas las posibilidades!"
105
+ '''
106
+ ]]
107
+ )
108
+
109
+ msg = gr.Textbox(placeholder="Escribe aquí", label='')
110
+ submit = gr.Button("Enviar", elem_id="enviar_button")
111
+ memory_state = gr.State(create_memory)
112
+
113
+ # Create the ChatOpenAI model with streaming enabled and our custom callback.
114
+ llm = ChatOpenAI(
115
+ temperature=0,
116
+ model_name='gpt-4o',
117
+ streaming=True,
118
+ callbacks=[stream_handler]
119
+ )
120
+
121
+ # --------------------------
122
+ # Generator function that runs the chain in a separate thread and polls the token queue.
123
+ # --------------------------
124
+ def user(query, chat_history, memory):
125
+ # Append the user's message with an empty bot response.
126
+ chat_history.append((query, ""))
127
+ # Immediately yield an update so the user's message appears.
128
+ yield "", chat_history, memory
129
+
130
+ # Container for the final chain result.
131
+ final_result = [None]
132
+
133
+ # Define a helper function to run the chain.
134
+ def run_chain():
135
+ result = pdf_qa(query, memory, llm)
136
+ final_result[0] = result
137
+ # Signal end-of-stream by putting a sentinel value.
138
+ stream_handler.token_queue.put(None)
139
+
140
+ # Run the chain in a separate thread.
141
+ thread = threading.Thread(target=run_chain)
142
+ thread.start()
143
+
144
+ # Poll the token queue for new tokens and yield updated chat history.
145
+ current_response = ""
146
+ while True:
147
+ try:
148
+ token = stream_handler.token_queue.get(timeout=0.1)
149
+ except queue.Empty:
150
+ token = None
151
+
152
+ # A None token is our signal for end-of-stream.
153
+ if token is None:
154
+ if not thread.is_alive():
155
+ break
156
+ else:
157
+ continue
158
+ current_response += token
159
+ chat_history[-1] = (query, current_response)
160
+ yield "", chat_history, memory
161
+
162
+ thread.join()
163
+ # Optionally, update the final answer if it differs from the streaming tokens.
164
+ if final_result[0] and "answer" in final_result[0]:
165
+ chat_history[-1] = (query, final_result[0]["answer"])
166
+ yield "", chat_history, memory
167
+
168
+ # Wire up the generator function to Gradio components with queue enabled.
169
+ submit.click(user, [msg, chatbot, memory_state], [msg, chatbot, memory_state], queue=True)
170
+ msg.submit(user, [msg, chatbot, memory_state], [msg, chatbot, memory_state], queue=True)
171
+
172
+ if __name__ == "__main__":
173
+ demo.queue().launch()