Undo changes in rag.py

appStore/rag.py

@@ -58,40 +58,27 @@ def run_query(context, label, model_sel_name):
     # Initialize the client, pointing it to one of the available models
     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
-    # )
-
-    # # Create an object to store the full chat completion
-    # completion_result = ChatCompletionResult()
-    # res_box = st.empty()
-
-    # # Iterate through the streamed output
-    # for chunk in chat_completion:
-    #     # Extract the object containing the text
-    #     if chunk.choices is not None:
-    #         chunk_message = chunk.choices[0].delta
-    #         if 'content' in chunk_message:
-    #             completion_result.add_content(chunk_message['content'])  # Store the message
-    #             # Add the latest text and merge it with all previous
-    #             result = completion_result.get_full_content()
-    #             res_box.success(result)  # Output to response text box
-
-    
-    # Use streaming text generation
-    response_stream = client.text_generation(prompt, stream=True, max_new_tokens=512)
+    # Instantiate ChatCompletion as a generator object (stream is set to True)
+    chat_completion = client.chat.completions.create(
+        messages=messages,
+        stream=True
+    )
 
+    # Create an object to store the full chat completion
     completion_result = ChatCompletionResult()
     res_box = st.empty()
 
-    for chunk in response_stream:
-        completion_result.add_content(chunk)
-        result = completion_result.get_full_content()
-        res_box.success(result)
+    # Iterate through the streamed output
+    for chunk in chat_completion:
+        # Extract the object containing the text
+        if chunk.choices is not None:
+            chunk_message = chunk.choices[0].delta
+            if 'content' in chunk_message:
+                completion_result.add_content(chunk_message['content'])  # Store the message
+                # Add the latest text and merge it with all previous
+                result = completion_result.get_full_content()
+                res_box.success(result)  # Output to response text box
 
-    return completion_result
 
-    # # Return the stored chat completion object for later use
-    # return completion_result
+    # Return the stored chat completion object for later use
+    return completion_result