achuthc1298 commited on
Commit
6603ddf
1 Parent(s): 2547926

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +1 -167
app.py CHANGED
@@ -1,168 +1,2 @@
1
- import gradio as gr
2
- import spaces
3
  import os
4
- from pathlib import Path
5
- from llama_index.core.tools import QueryEngineTool
6
- from llama_index.core import VectorStoreIndex
7
- from llama_index.core import Settings
8
- from llama_index.core import SimpleDirectoryReader
9
- from llama_index.llms.groq import Groq
10
- from llama_index.embeddings.huggingface import HuggingFaceEmbedding
11
- from typing import Tuple
12
- from llama_index.core import StorageContext, load_index_from_storage
13
- from llama_index.core.objects import ObjectIndex
14
- from llama_index.core.agent import ReActAgent
15
- import torch
16
- import sys
17
- import io
18
-
19
- @spaces.GPU()
20
- def create_doc_tools(document_fp: str, doc_name: str, verbose: bool = True) -> Tuple[QueryEngineTool,]:
21
- documents = SimpleDirectoryReader(input_files=[document_fp]).load_data()
22
-
23
- llm = Groq(model="mixtral-8x7b-32768")
24
- device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
25
- #local_model_path = "/home/user/app/sentence-transformers--all-mpnet-base-v2"
26
- #embed_model = HuggingFaceEmbedding(model_name=local_model_path, device=device)
27
- embed_model = HuggingFaceEmbedding(model_name="sentence-transformers/all-mpnet-base-v2", device=device)
28
- Settings.llm = llm
29
- Settings.embed_model = embed_model
30
-
31
- load_dir_path = f"/home/user/app/agentic_index_st/{doc_name}"
32
- storage_context = StorageContext.from_defaults(persist_dir=load_dir_path)
33
- vector_index = load_index_from_storage(storage_context)
34
- vector_query_engine = vector_index.as_query_engine()
35
-
36
- vector_tool = QueryEngineTool.from_defaults(
37
- name=f"{doc_name}_vector_query_engine_tool",
38
- query_engine=vector_query_engine,
39
- description=f"Useful for retrieving specific context from the {doc_name}.",
40
- )
41
-
42
- return vector_tool
43
-
44
- def find_tex_files(directory: str):
45
- tex_files = []
46
- for root, dirs, files in os.walk(directory):
47
- for file in files:
48
- if file.endswith(('.tex', '.txt')):
49
- file_path = os.path.abspath(os.path.join(root, file))
50
- tex_files.append(file_path)
51
- tex_files.sort()
52
- return tex_files
53
-
54
- def initialize_agent(apikey):
55
- os.environ["GROQ_API_KEY"] = apikey
56
-
57
- llm = Groq(model="mixtral-8x7b-32768")
58
-
59
- try:
60
- directory = '/home/user/app/rag_docs_final_review_tex_merged'
61
- tex_files = find_tex_files(directory)
62
-
63
- paper_to_tools_dict = {}
64
- for paper in tex_files:
65
- path = Path(paper)
66
- vector_tool = create_doc_tools(doc_name=path.stem, document_fp=path)
67
- paper_to_tools_dict[path.stem] = [vector_tool]
68
-
69
- initial_tools = [t for paper in tex_files for t in paper_to_tools_dict[Path(paper).stem]]
70
-
71
- obj_index = ObjectIndex.from_objects(
72
- initial_tools,
73
- index_cls=VectorStoreIndex,
74
- )
75
-
76
- obj_retriever = obj_index.as_retriever(similarity_top_k=6)
77
-
78
- context = """You are an agent designed to answer scientific queries over a set of given documents.
79
- Please always use the tools provided to answer a question. Do not rely on prior knowledge.
80
- """
81
-
82
- agent = ReActAgent.from_tools(
83
- tool_retriever=obj_retriever,
84
- llm=llm,
85
- verbose=True,
86
- context=context
87
- )
88
- return agent
89
- except Exception as e:
90
- return f"Error: {str(e)}"
91
-
92
- def chat_with_agent(prompt, agent, verbose_toggle):
93
- try:
94
- # Check if agent initialization was successful
95
- if isinstance(agent, str) and agent.startswith("Error:"):
96
- return agent # Return the error message if initialization failed
97
-
98
- # Redirect stdout
99
- original_stdout = sys.stdout
100
- sys.stdout = io.StringIO()
101
-
102
- # query the agent
103
- response = agent.query(prompt)
104
-
105
- # Get the captured output and restore stdout
106
- output = sys.stdout.getvalue()
107
- sys.stdout = original_stdout
108
-
109
- # format the received verbose output
110
- verbose = ''
111
- for output_string in output.split('==='):
112
- verbose += output_string
113
- verbose += '\n'
114
-
115
- # assistant response
116
- msg = f'{verbose}' if verbose_toggle else f'{response.response[10:]}'
117
-
118
- return msg
119
- except Exception as e:
120
- return str(e)
121
-
122
- def main():
123
- agent = None
124
-
125
- def set_apikey(apikey):
126
- nonlocal agent
127
- agent = initialize_agent(apikey)
128
- if isinstance(agent, str) and agent.startswith("Error:"):
129
- return agent # Return the error message if initialization failed
130
- return "API Key Set. You may start asking questions now."
131
-
132
- def reset_chat():
133
- nonlocal agent
134
- agent = None
135
- return "Chat reset. You may start asking questions now."
136
-
137
- def chat_function(prompt, apikey, verbose_toggle):
138
- nonlocal agent
139
- if not agent:
140
- api_response = set_apikey(apikey)
141
- if isinstance(agent, str) and agent.startswith("Error:"):
142
- return api_response # Return the error message if initialization failed
143
- return chat_with_agent(prompt, agent, verbose_toggle)
144
-
145
- with gr.Blocks() as demo:
146
- gr.Markdown("# AMGPT, Powered by LlamaIndex")
147
-
148
- with gr.Row():
149
- apikey = gr.Textbox(label="Enter your Groq API Key", type="password")
150
- set_apikey_button = gr.Button("Set API Key")
151
-
152
- set_apikey_button.click(set_apikey, inputs=apikey, outputs=None)
153
-
154
- with gr.Row():
155
- verbose_toggle = gr.Checkbox(label="Verbose", value=True)
156
- reset = gr.Button("Reset Chat")
157
-
158
- reset.click(reset_chat, outputs=None)
159
-
160
- prompt = gr.Textbox(label="Ask a question")
161
- output = gr.Textbox(label="Response")
162
-
163
- prompt.submit(chat_function, inputs=[prompt, apikey, verbose_toggle], outputs=output)
164
-
165
- demo.launch()
166
-
167
- if __name__ == "__main__":
168
- main()
 
 
 
1
  import os
2
+ print(os.getcwd())