Update appStore/rag.py

appStore/rag.py

@@ -10,11 +10,12 @@ 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. \
+  If there is no mention of "+label+" in the context, return: 'No clear references to vulnerability of "+label+" found'.  \
   Do not include an introduction sentence, just the bullet points as per below. \
   Formatting example: \
     - Bullet point 1 \
@@ -32,7 +33,7 @@ def get_prompt(context, label):
 #     return openai.ChatCompletion.create(**kwargs)
 
 # construct query, send to HF API and process response
-def run_query(context, label):
+def run_query(context, label, model_sel_name):
     '''
     For non-streamed completion, enable the following 2 lines and comment out the code below
     '''
@@ -40,29 +41,29 @@ 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.1-405B-Instruct", token = hf_token)
-
+    client = InferenceClient(model_sel_name, token = hf_token)
+    
+    # instantiate ChatCompletion as a generator object (stream is set to True)
     chat_completion = client.chat.completions.create(
         messages=messages,
         stream=True
     )
+    # chat_completion = completion_with_backoff(messages=messages, stream=True)
 
     # iterate through the streamed output
     report = []
     res_box = st.empty()
     for chunk in chat_completion:
         # extract the object containing the text (totally different structure when streaming)
-        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
-            # add the latest text and merge it with all previous
-            result = "".join(report).strip()
-            # res_box.success(result) # output to response text box
-            res_box.success(result)
-
-            if chunk.choices[0].finish_reason != None:
-                break
+        if chunk.choices is not None: # sometimes returns None - probably the prompt needs work
+            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
+                # add the latest text and merge it with all previous
+                result = "".join(report).strip()
+                # res_box.success(result) # output to response text box
+                res_box.success(result)