Update appStore/rag.py

appStore/rag.py

@@ -11,17 +11,19 @@ from huggingface_hub import InferenceClient
 # Get openai API key
 hf_token = os.environ["HF_API_KEY"]
 
-
 # define a special function for putting the prompt together (as we can't use haystack)
 def get_prompt(context, label):
   base_prompt="Summarize the following context efficiently in bullet points, the less the better - but keep concrete goals. \
   Summarize only elements of the context that address vulnerability of "+label+" to climate change. \
   If there is no mention of "+label+" in the context, return nothing. \
+  Do not include an introduction sentence, just the bullet points as per below. \
   Formatting example: \
     - Bullet point 1 \
-    - Bullet point 2 "
+    - Bullet point 2 \
+  "
+
   prompt = base_prompt+"; Context: "+context+"; Answer:"
-  
+
   return prompt
 
 
@@ -39,16 +41,19 @@ def run_query(context, label):
     messages = [{"role": "system", "content": chatbot_role},{"role": "user", "content": get_prompt(context, label)}]
     
     # Initialize the client, pointing it to one of the available models
-    client = InferenceClient("meta-llama/Meta-Llama-3-8B-Instruct", token = hf_token)
+    client = InferenceClient("meta-llama/Meta-Llama-3.1-405B-Instruct", token = hf_token)
+
+    chat_completion = client.chat.completions.create(
+        messages=messages,
+        stream=True
+    )
 
-    # instantiate ChatCompletion as a generator object (stream is set to True)
-    # response = completion_with_backoff(model=model_select, messages=[{"role": "user", "content": get_prompt(context, label)}], stream=True)
     # iterate through the streamed output
     report = []
     res_box = st.empty()
-    for chunk in client.chat_completion(messages, stream=True):
+    for chunk in chat_completion:
         # extract the object containing the text (totally different structure when streaming)
-        chunk_message = chunk['choices'][0]['delta']
+        chunk_message = chunk.choices[0].delta
         # test to make sure there is text in the object (some don't have)
         if 'content' in chunk_message:
             report.append(chunk_message['content']) # extract the message
@@ -63,6 +68,3 @@ def run_query(context, label):
 
 
 
-
-
-