Update app.py

app.py

@@ -8,31 +8,35 @@ from transformers import pipeline
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
 from helpers.foundation_models import *
+import requests
+
 
 OPENAI_API_KEY = os.environ["OPENAI_API_KEY"]
 SERPAPI_API_KEY = os.environ["SERPAPI_API_KEY"]
 openai_client = openai.OpenAI(api_key=OPENAI_API_KEY)
 
 
-# tokenizer = AutoTokenizer.from_pretrained("eagle0504/llama-2-7b-miniguanaco")
-# model = AutoModelForCausalLM.from_pretrained("eagle0504/llama-2-7b-miniguanaco")
-
+API_URL = "https://sks7h7h5qkhoxwxo.us-east-1.aws.endpoints.huggingface.cloud"
+headers = {
+	"Accept" : "application/json",
+	"Content-Type": "application/json" 
+}
 
-# def generate_response_from_llama2(query):
 
-#     # Tokenize the input text
-#     input_ids = tokenizer.encode(query, return_tensors="pt")
+def query(payload):
+	response = requests.post(API_URL, headers=headers, json=payload)
+	return response.json()
 
-#     # Generate a response
-#     # Adjust the parameters like max_length according to your needs
-#     output = model.generate(input_ids, max_length=50, num_return_sequences=1, temperature=0.7)
 
-#     # Decode the output to human-readable text
-#     generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
+def llama2_7b_ysa(prompt: str) -> str:
+    output = query({
+        "inputs": prompt,
+        "parameters": {}
+    })
 
-#     # output
-#     return generated_text
+    response = output[0]['generated_text']
 
+    return response
 
 # Initialize chat history
 if "messages" not in st.session_state:
@@ -59,7 +63,7 @@ with st.expander("Instructions"):
 
 option = st.sidebar.selectbox(
     "Which task do you want to do?",
-    ("Sentiment Analysis", "Medical Summarization", "ChatGPT", "ChatGPT (with Google)"),
+    ("Sentiment Analysis", "Medical Summarization", "Llama2 on YSA", "ChatGPT", "ChatGPT (with Google)"),
 )
 
 
@@ -87,7 +91,7 @@ if prompt := st.chat_input("What is up?"):
             pipe_sentiment_analysis = pipeline("sentiment-analysis")
             if prompt:
                 out = pipe_sentiment_analysis(prompt)
-                doc = f"""
+                final_response = f"""
                     Prompt: {prompt}
                     Sentiment: {out[0]["label"]}
                     Score: {out[0]["score"]}
@@ -98,15 +102,22 @@ if prompt := st.chat_input("What is up?"):
             )
             if prompt:
                 out = pipe_summarization(prompt)
-                doc = out[0]["summary_text"]
-        # elif option == "Llama2":
-        #     if prompt:
-        #         out = generate_response_from_llama2(query=prompt)
-        #         doc = out
+                final_response = out[0]["summary_text"]
+        elif option == "Llama2 on YSA":
+            if prompt:
+                out = llama2_7b_ysa(query=prompt)
+                engineered_prompt = f"""
+                    The user asked the question: {prompt}
+
+                    We have found relevant content: {out}
+
+                    Answer the user question based on the above content in paragraphs.
+                """
+                final_response = call_chatgpt(query=engineered_prompt)
         elif option == "ChatGPT":
             if prompt:
                 out = call_chatgpt(query=prompt)
-                doc = out
+                final_response = out
         elif option == "ChatGPT (with Google)":
             if prompt:
                 ans_langchain = call_langchain(prompt)
@@ -116,11 +127,11 @@ if prompt := st.chat_input("What is up?"):
                     Answer the user question: {prompt}
                 """
                 out = call_chatgpt(query=prompt)
-                doc = out
+                final_response = out
         else:
-            doc = ""
+            final_response = ""
 
-    response = f"{doc}"
+    response = f"{final_response}"
     # Display assistant response in chat message container
     with st.chat_message("assistant"):
         st.markdown(response)