Update app.py

app.py

@@ -1,34 +1,24 @@
-import gradio as gr
 import os
-import spaces
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
-from threading import Thread
-
-# Set an environment variable
-HF_TOKEN = os.environ.get("HF_TOKEN", None)
-
-SYSTEM_PROMPT = "You are an AI Healthcare assistant. Provide clear, empathetic, and accurate information about healthcare topics. Do not state It's always best to consult with a healthcare provider, as you are simulating the role of a healthcare provider. A disclaimer has been placed elsewhere for the user that informs them to be mindful about your replies, so your job is to provide them with as much information as possible and let them make decisions themselves. Only refer them to an outside provider for physical resources, such as prescriptions, tests, bloodwork, x-rays, ultrasounds, CT scans, MRIs, and other similar things, since you are virtual."
+import gradio as gr
+from openai import OpenAI
 
+# Configure the OpenAI client with your custom API endpoint and API key.
+client = OpenAI(base_url="http://home.mayhew.cloud:1234/v1", api_key="lm-studio")
 
+# UI text and styling
+SYSTEM_PROMPT = "You are an assistant."
 DESCRIPTION = '''
 <div>
-<h1 style="text-align: center;">HealthAssistant</h1>
+  <h1 style="text-align: center;">HealthAssistant</h1>
 </div>
 '''
-
-LICENSE = """
-<p>
-This Health Assistant is designed to provide helpful healthcare information; however, it may make mistakes and is not designed to replace professional medical care. It is not intended to diagnose any condition or disease. Always consult with a qualified healthcare provider for any medical concerns.\n\nI hereby confirm that I am at least 18 years of age (or accompanied by a legal guardian who is at least 18 years old), understand that the information provided by this service is for informational purposes only and is not intended to diagnose or treat any medical condition, and acknowledge that I am solely responsible for verifying any information provided.
-</p>
-"""
-
+LICENSE = "<p></p>"
 PLACEHOLDER = """
 <div style="padding: 30px; text-align: center; display: flex; flex-direction: column; align-items: center;">
    <h1 style="font-size: 28px; margin-bottom: 2px; opacity: 0.55;">The "Doctor" is in.</h1>
    <p style="font-size: 18px; margin-bottom: 2px; opacity: 0.65;">Available 1:00pm - 5:00pm EST</p>
 </div>
 """
-
 css = """
 h1 {
   text-align: center;
@@ -43,116 +33,95 @@ h1 {
 }
 """
 
-# Load the tokenizer and model
-tokenizer = AutoTokenizer.from_pretrained("reedmayhew/HealthCare-Reasoning-Assistant-Llama-3.1-8B-HF", device_map="cuda")
-model = AutoModelForCausalLM.from_pretrained("reedmayhew/HealthCare-Reasoning-Assistant-Llama-3.1-8B-HF", device_map="cuda")
-
-terminators = [
-    tokenizer.eos_token_id,
-    tokenizer.convert_tokens_to_ids("<|eot_id|>")
-]
-
-@spaces.GPU(duration=60)
-def chat_llama3_8b(message: str, 
-                    history: list, 
-                    temperature: float, 
-                    max_new_tokens: int
-                   ) -> str:
+def chat_with_openai(message: str, history: list, temperature: float, max_new_tokens: int):
     """
-    Generate a streaming response using the llama3-8b model.
+    Call the OpenAI ChatCompletion endpoint using the new client and yield streaming responses.
+    Implements <think> logic:
+      - The assistant is forced to begin its answer with "<think> ".
+      - We then wait until a closing "</think>" marker is received.
+      - Only text after "</think>" is displayed as the final answer.
+    
     Args:
-        message (str): The input message.
-        history (list): The conversation history used by ChatInterface.
-        temperature (float): The temperature for generating the response.
-        max_new_tokens (int): The maximum number of new tokens to generate.
-    Returns:
-        str: The generated response.
-    """
+        message (str): The latest user message.
+        history (list): Conversation history as a list of (user, assistant) tuples.
+        temperature (float): Sampling temperature.
+        max_new_tokens (int): Maximum tokens to generate.
     
+    Yields:
+        str: Partial cumulative output from the assistant.
+    """
     conversation = []
-
-    # Add system prompt only if this is a new conversation (i.e., history is empty)
     if not history:
+        # Add a system prompt and initial assistant confirmation.
         conversation.append({"role": "system", "content": SYSTEM_PROMPT})
-        conversation.append({"role": "assistant", "content": "Understood! I will act as the user's healthcare provider, not hesitate with providing detailed responses, or refer them to another healthcare provider unless they ask for one or need a physical resource, as you have already given them a disclaimer."})
-
-    
-    for user, assistant in history:
-        conversation.extend([
-            {"role": "user", "content": user}, 
-            {"role": "assistant", "content": assistant}
-        ])
-    
-    # Ensure the model starts with "<think>"
+        conversation.append({"role": "assistant", "content": "Understood!"})
+    for user_msg, assistant_msg in history:
+        conversation.append({"role": "user", "content": user_msg})
+        conversation.append({"role": "assistant", "content": assistant_msg})
     conversation.append({"role": "user", "content": message})
-    conversation.append({"role": "assistant", "content": "<think> "})  # Force <think> at start
+    # Force the model to begin its answer with a "<think>" block.
+    conversation.append({"role": "assistant", "content": "<think> "})
 
-    input_ids = tokenizer.apply_chat_template(conversation, return_tensors="pt").to(model.device)
-    
-    streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True)
+    full_response = ""   # Stores the raw assistant response (including the <think> block).
+    buffer = ""          # Accumulates tokens until we detect the closing </think>.
+    display_text = ""    # Holds text to display (only text after </think>).
+    think_detected = False
+
+    # Immediately yield a "thinking" status message.
+    yield "A.I. Healthcare is Thinking! Please wait, your response will output shortly...\n\n"
 
-    generate_kwargs = dict(
-        input_ids=input_ids,
-        streamer=streamer,
-        max_new_tokens=max_new_tokens,
-        do_sample=True,
+    # Call the API with streaming enabled.
+    response = client.chat.completions.create(
+        model="model-identifier",  # Replace with your actual model identifier.
+        messages=conversation,
         temperature=temperature,
-        eos_token_id=terminators,
+        max_tokens=max_new_tokens,
+        stream=True,
     )
-    
-    if temperature == 0:
-        generate_kwargs['do_sample'] = False
-        
-    t = Thread(target=model.generate, kwargs=generate_kwargs)
-    t.start()
-
-    outputs = []
-    buffer = ""
-    think_detected = False
-    thinking_message_sent = False
-    full_response = ""  # Store the full assistant response
 
-    for text in streamer:
-        buffer += text
-        full_response += text  # Store raw assistant response (includes <think>)
+    # Process streaming responses.
+    for chunk in response:
+        # Extract the new token text from the chunk.
+        delta = chunk.choices[0].delta
+        token_text = delta.content or ""
+        full_response += token_text
 
-        # Send the "thinking" message once text starts generating
-        if not thinking_message_sent:
-            thinking_message_sent = True
-            yield "A.I. Healthcare is Thinking! Please wait, your response will output shortly...\n\n"
-
-        # Wait until </think> is detected before streaming output
         if not think_detected:
+            # Accumulate tokens until we see the closing </think> marker.
+            buffer += token_text
             if "</think>" in buffer:
                 think_detected = True
-                buffer = buffer.split("</think>", 1)[1]  # Remove <think> section
+                # Discard everything up to and including the "</think>" marker.
+                display_text = buffer.split("</think>", 1)[1]
+                yield display_text
         else:
-            outputs.append(text)
-            yield "".join(outputs)
+            display_text += token_text
+            yield display_text
 
-    # Store the full response (including <think>) in history, but only show the user the cleaned response
-    history.append((message, full_response))  # Full assistant response saved for context
+    # Append the full (raw) response, including the <think> section, to the conversation history.
+    history.append((message, full_response))
 
-# Gradio block
+# Create the Chatbot component.
 chatbot = gr.Chatbot(height=450, placeholder=PLACEHOLDER, label='HealthAssistant')
 
-with gr.Blocks(fill_height=True, css=css) as demo:
-    
+# Build the Gradio interface.
+with gr.Blocks(css=css) as demo:
     gr.Markdown(DESCRIPTION)
+    
     gr.ChatInterface(
-        fn=chat_llama3_8b,
+        fn=chat_with_openai,
         chatbot=chatbot,
         fill_height=True,
         additional_inputs_accordion=gr.Accordion(label="⚙️ Parameters", open=False, render=False),
         additional_inputs=[
-            gr.Slider(minimum=0.6, maximum=0.6, step=0.1, value=0.6, label="Temperature", render=False),
+            gr.Slider(minimum=0.0, maximum=1.0, step=0.1, value=0.6, label="Temperature", render=False),
             gr.Slider(minimum=1024, maximum=4096, step=128, value=2048, label="Max new tokens", render=False),
         ],
         examples=[
-            ['What is PrEP, and do I need it?'],
-            ['What medications help manage being undetectable with HIV?'],
-            ['How do I know if an abortion is the right option?'],
-            ['How can I access birth-control in states where it is regulated?'],
+            ['What is, and do I need it?'],
+            ['What medications help manage being invisible?'],
+            ['How do I know if a clown is the right option?'],
+            ['How can I access music in states where it is regulated?'],
         ],
         cache_examples=False,
     )