Update app.py

app.py

@@ -14,6 +14,7 @@ from llama_index.retrievers.bm25 import BM25Retriever
 from llama_index.core.retrievers import BaseRetriever
 from llama_index.core.node_parser import SentenceSplitter
 from llama_index.embeddings.openai import OpenAIEmbedding
+from llama_index.llms.mistralai import MistralAI
 from llmlingua import PromptCompressor
 from rouge_score import rouge_scorer
 from semantic_text_similarity.models import WebBertSimilarity
@@ -27,6 +28,13 @@ nest_asyncio.apply()
 # openai.api_key = key 
 # os.environ["OPENAI_API_KEY"] = key
 
+# key = os.getenv('MISTRAL_API_KEY')
+# os.environ["MISTRAL_API_KEY"] = key
+
+# Anthropic credentials
+key = os.getenv('CLAUDE_API_KEY')
+os.environ["ANTHROPIC_API_KEY"] = key
+
 # Streamlit UI
 st.title("Prompt Optimization for a Policy Bot")
 
@@ -38,6 +46,7 @@ if uploaded_files:
             f.write(uploaded_file.getbuffer())
         reader = SimpleDirectoryReader(input_files=[f"./data/{uploaded_file.name}"])
         documents = reader.load_data()
+        st.write(documents)
         st.success("File uploaded...")
 
         # # Indexing
@@ -103,7 +112,8 @@ if uploaded_files:
         hybrid_retriever = HybridRetriever(vector_retriever, bm25_retriever)
 
         # Generation
-        model = "gpt-3.5-turbo"
+        # model = "gpt-3.5-turbo"
+        model = "claude-3-opus-20240229"
 
         # def get_context(query):
         #     contexts = kg_retriever.retrieve(query)
@@ -145,6 +155,10 @@ if uploaded_files:
             with st.chat_message(message["role"]):
                 st.markdown(message["content"])
 
+        # Summarize 
+        full_prompt = "\n\n".join([context + prompt])
+        orig_res = res(full_prompt)
+
         # Accept user input
         if prompt := st.chat_input("Enter your query:"):
             st.success("Fetching info...")
@@ -158,6 +172,11 @@ if uploaded_files:
             context_list = get_context(prompt)
             context = " ".join(context_list)
 
+            # # Summarize 
+            # full_prompt = "\n\n".join([context + prompt])
+            # orig_res = res(full_prompt)
+            
+
 
             # Original prompt response
             full_prompt = "\n\n".join([context + prompt])
@@ -168,66 +187,66 @@ if uploaded_files:
             with st.chat_message("assistant"):
                 st.markdown(orig_res[3])
 
-            # Compressed Response
-            st.session_state.messages.append({"role": "assistant", "content": "Generating Optimized prompt response..."})
-            st.success("Generating Optimized prompt response...")
-
-            llm_lingua = PromptCompressor(
-            model_name="microsoft/llmlingua-2-xlm-roberta-large-meetingbank",
-            use_llmlingua2=True, device_map="cpu"
-            )
-
-            def prompt_compression(context, rate=0.5):
-                compressed_context = llm_lingua.compress_prompt(
-                    context,
-                    rate=rate,
-                    force_tokens=["!", ".", "?", "\n"],
-                    drop_consecutive=True,
-                )
-                return compressed_context
-            compressed_context = prompt_compression(context)
-            full_opt_prompt = "\n\n".join([compressed_context['compressed_prompt'] + prompt])
-            compressed_res = res(full_opt_prompt)
-            st.session_state.messages.append({"role": "assistant", "content": compressed_res[3]})
-            with st.chat_message("assistant"):
-                st.markdown(compressed_res[3])
-
-            # Save token summary and evaluation details to session state
-            scorer = rouge_scorer.RougeScorer(['rougeL'], use_stemmer=True)
-            scores = scorer.score(compressed_res[3],orig_res[3])
-            webert_model = WebBertSimilarity(device='cpu')
-            similarity_score = webert_model.predict([(compressed_res[3], orig_res[3])])[0] / 5 * 100
+            # # Compressed Response
+            # st.session_state.messages.append({"role": "assistant", "content": "Generating Optimized prompt response..."})
+            # st.success("Generating Optimized prompt response...")
+
+            # llm_lingua = PromptCompressor(
+            # model_name="microsoft/llmlingua-2-xlm-roberta-large-meetingbank",
+            # use_llmlingua2=True, device_map="cpu"
+            # )
+
+            # def prompt_compression(context, rate=0.5):
+            #     compressed_context = llm_lingua.compress_prompt(
+            #         context,
+            #         rate=rate,
+            #         force_tokens=["!", ".", "?", "\n"],
+            #         drop_consecutive=True,
+            #     )
+            #     return compressed_context
+            # compressed_context = prompt_compression(context)
+            # full_opt_prompt = "\n\n".join([compressed_context['compressed_prompt'] + prompt])
+            # compressed_res = res(full_opt_prompt)
+            # st.session_state.messages.append({"role": "assistant", "content": compressed_res[3]})
+            # with st.chat_message("assistant"):
+            #     st.markdown(compressed_res[3])
+
+            # # Save token summary and evaluation details to session state
+            # scorer = rouge_scorer.RougeScorer(['rougeL'], use_stemmer=True)
+            # scores = scorer.score(compressed_res[3],orig_res[3])
+            # webert_model = WebBertSimilarity(device='cpu')
+            # similarity_score = webert_model.predict([(compressed_res[3], orig_res[3])])[0] / 5 * 100
          
 
-            # Display token summary
-            st.session_state.messages.append({"role": "assistant", "content": "Token Length Summary..."})
-            st.success('Token Length Summary...')
-            st.session_state.messages.append({"role": "assistant", "content": f"Original Prompt has {orig_res[0]} tokens"})
-            st.write(f"Original Prompt has {orig_res[0]} tokens")
-            st.session_state.messages.append({"role": "assistant", "content": f"Optimized Prompt has {compressed_res[0]} tokens"})
-            st.write(f"Optimized Prompt has {compressed_res[0]} tokens")
-
-            st.session_state.messages.append({"role": "assistant", "content": "Comparing Original and Optimized Prompt Response..."})
-            st.success("Comparing Original and Optimized Prompt Response...")
-            st.session_state.messages.append({"role": "assistant", "content": f"Rouge Score : {scores['rougeL'].fmeasure * 100}"})
-            st.write(f"Rouge Score : {scores['rougeL'].fmeasure * 100}")
-            st.session_state.messages.append({"role": "assistant", "content": f"Semantic Text Similarity Score : {similarity_score}"})
-            st.write(f"Semantic Text Similarity Score : {similarity_score}")
-
-            st.write(" ")
-            # origin_tokens = compressed_context['origin_tokens']
-            # compressed_tokens = compressed_context['compressed_tokens']
-            origin_tokens = orig_res[0]
-            compressed_tokens = compressed_res[0]
-            gpt_saving = (origin_tokens - compressed_tokens) * 0.06 / 1000
-            claude_saving = (origin_tokens - compressed_tokens) * 0.015 / 1000
-            mistral_saving = (origin_tokens - compressed_tokens) * 0.004 / 1000
-            # st.session_state.messages.append({"role": "assistant", "content": f"""The optimized prompt has saved ${gpt_saving:.4f} in GPT4, ${mistral_saving:.4f} in Mistral"""})
-            # st.success(f"""The optimized prompt has saved ${gpt_saving:.4f} in GPT4, ${mistral_saving:.4f} in Mistral""")
-            st.session_state.messages.append({"role": "assistant", "content": f"The optimized prompt has ${gpt_saving:.4f} saved in GPT-4."})
-            st.success(f"The optimized prompt has ${gpt_saving:.4f} saved in GPT-4.")
-
-            st.success("Downloading Optimized Prompt...")
-            st.download_button(label = "Download Optimized Prompt", 
-                               data = full_opt_prompt, file_name='./data/optimized_prompt.txt')
+            # # Display token summary
+            # st.session_state.messages.append({"role": "assistant", "content": "Token Length Summary..."})
+            # st.success('Token Length Summary...')
+            # st.session_state.messages.append({"role": "assistant", "content": f"Original Prompt has {orig_res[0]} tokens"})
+            # st.write(f"Original Prompt has {orig_res[0]} tokens")
+            # st.session_state.messages.append({"role": "assistant", "content": f"Optimized Prompt has {compressed_res[0]} tokens"})
+            # st.write(f"Optimized Prompt has {compressed_res[0]} tokens")
+
+            # st.session_state.messages.append({"role": "assistant", "content": "Comparing Original and Optimized Prompt Response..."})
+            # st.success("Comparing Original and Optimized Prompt Response...")
+            # st.session_state.messages.append({"role": "assistant", "content": f"Rouge Score : {scores['rougeL'].fmeasure * 100}"})
+            # st.write(f"Rouge Score : {scores['rougeL'].fmeasure * 100}")
+            # st.session_state.messages.append({"role": "assistant", "content": f"Semantic Text Similarity Score : {similarity_score}"})
+            # st.write(f"Semantic Text Similarity Score : {similarity_score}")
+
+            # st.write(" ")
+            # # origin_tokens = compressed_context['origin_tokens']
+            # # compressed_tokens = compressed_context['compressed_tokens']
+            # origin_tokens = orig_res[0]
+            # compressed_tokens = compressed_res[0]
+            # gpt_saving = (origin_tokens - compressed_tokens) * 0.06 / 1000
+            # claude_saving = (origin_tokens - compressed_tokens) * 0.015 / 1000
+            # mistral_saving = (origin_tokens - compressed_tokens) * 0.004 / 1000
+            # # st.session_state.messages.append({"role": "assistant", "content": f"""The optimized prompt has saved ${gpt_saving:.4f} in GPT4, ${mistral_saving:.4f} in Mistral"""})
+            # # st.success(f"""The optimized prompt has saved ${gpt_saving:.4f} in GPT4, ${mistral_saving:.4f} in Mistral""")
+            # st.session_state.messages.append({"role": "assistant", "content": f"The optimized prompt has ${gpt_saving:.4f} saved in GPT-4."})
+            # st.success(f"The optimized prompt has ${gpt_saving:.4f} saved in GPT-4.")
+
+            # st.success("Downloading Optimized Prompt...")
+            # st.download_button(label = "Download Optimized Prompt", 
+            #                    data = full_opt_prompt, file_name='./data/optimized_prompt.txt')