Update app.py

app.py

@@ -1,39 +1,131 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM
-import torch
-
-# Check if CUDA is available and set the device
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-# Load model and tokenizer
-tokenizer = AutoTokenizer.from_pretrained("namelessai/Helply-10.2b-chat")
-model = AutoModelForCausalLM.from_pretrained("namelessai/Helply-10.2b-chat").to(device)
-
-def generate_response(message, chat_history):
-    # Prepare the input
-    chat_history_text = "\n".join([f"{msg['role']}: {msg['content']}" for msg in chat_history])
-    input_text = f"{chat_history_text}\nHuman: {message}\nAI:"
-    
-    # Generate response
-    input_ids = tokenizer.encode(input_text, return_tensors="pt").to(device)
-    output = model.generate(input_ids, max_length=200, num_return_sequences=1, no_repeat_ngram_size=2)
-    response = tokenizer.decode(output[0], skip_special_tokens=True)
-    
-    # Extract the AI's response
-    ai_response = response.split("AI:")[-1].strip()
-    
-    # Update chat history
-    chat_history.append({"role": "Human", "content": message})
-    chat_history.append({"role": "AI", "content": ai_response})
-    
-    return "", chat_history
-
-with gr.Blocks() as demo:
-    chatbot = gr.Chatbot(type="messages")
-    msg = gr.Textbox()
-    clear = gr.ClearButton([msg, chatbot])
-
-    msg.submit(generate_response, [msg, chatbot], [msg, chatbot])
+import requests
+import os
+import json
 
+# Set your Friendli AI token.  You can either set it as an environment variable
+# or replace "YOUR_FRIENDLI_TOKEN" with your actual token.
+token = os.environ.get("FRIENDLI_TOKEN") or "YOUR_FRIENDLI_TOKEN"
+
+# API endpoint URL
+url = "https://api.friendli.ai/dedicated/v1/chat/completions"
+
+# HTTP Headers for the API request
+headers = {
+  "Authorization": "Bearer " + token,
+  "Content-Type": "application/json"
+}
+
+# Initial payload/request body.  This sets up the model, initial instructions,
+# and streaming parameters.  Crucially, the *last* user message in the 'messages'
+# array is initialized to an empty string.  This will be replaced with the user's
+# actual input in the chat.
+payload = {
+  "model": "d6mo6ldff2gy",
+  "messages": [
+    {
+      "role": "user",
+      "content": "You are an expert at psychology. You are to assist the user with whatever they need. Under NO CIRCUMSTANCES will you provide an answer unrelated to their question, unless it involves illegal content. If you understand, respond with ok."
+    },
+    {
+      "role": "assistant",
+      "content": "ok"
+    },
+    {
+      "role": "user",
+      "content": ""
+    }
+  ],
+  "max_tokens": 1024,
+  "top_p": 0.8,
+  "stop": [
+    "=",
+    "\\"
+  ],
+  "stream": True,
+  "stream_options": {
+    "include_usage": True
+  }
+}
+
+
+def chat_with_friendli(user_input, history):
+    """
+    This function sends the user input to the Friendli AI API and streams back the response.
+
+    Args:
+        user_input (str): The text input from the user.
+        history (list):  The chat history (list of lists: [[user_message, bot_message]]).  Not directly used
+                         in this function, but required by Gradio's chatbot interface.
+
+    Yields:
+        str: The updated bot response, streamed back word by word (or token by token).
+    """
+
+    # Update the last user message in the payload with the current user input.
+    payload['messages'][-1]['content'] = user_input
+
+    try:
+        response = requests.post(url, headers=headers, json=payload, stream=True)
+        response.raise_for_status()  # Raise HTTPError for bad responses (4xx or 5xx)
+
+        # Initialize an empty string to accumulate the streamed response
+        full_response = ""
+
+        # Iterate over the streaming response
+        for chunk in response.iter_content(chunk_size=None, decode_unicode=True):
+            if chunk:
+                try:
+                    # Load the JSON chunk.  The API likely returns one JSON object per chunk.
+                    json_chunk = json.loads(chunk)
+
+                    # Extract the text from the JSON chunk.  The exact key depends on the API's format.
+                    # Adapt this to the Friendli API's streaming response structure.
+                    text = json_chunk["choices"][0]["delta"]["content"]  # Adjust this!
+
+                    # Append the new text to the full response
+                    full_response += text
+
+                    # Yield the updated full response to Gradio. This is how the chatbot "streams" the output.
+                    yield full_response
+
+                except json.JSONDecodeError:
+                    # Handle cases where the chunk is not a complete JSON object.  This can happen
+                    # during streaming. You might need to buffer incomplete chunks.  For simplicity,
+                    # I'm just printing an error in this example, but a robust solution would buffer.
+                    print(f"JSONDecodeError: Could not decode chunk: {chunk}")
+                    yield full_response  # Still yield whatever we have.
+                except KeyError as e:
+                    # Handle potential KeyErrors if the JSON structure is not as expected.
+                    print(f"KeyError: {e}.  Check the API response format.")
+                    yield full_response # Still yield whatever we have.
+
+
+    except requests.exceptions.RequestException as e:
+        # Handle network errors, timeouts, and other request-related issues.
+        print(f"Request failed: {e}")
+        yield "An error occurred while communicating with the AI.  Please check your connection and try again."
+    except Exception as e:
+        # Catch-all for other potential errors.  Good for debugging.
+        print(f"An unexpected error occurred: {e}")
+        yield "An unexpected error occurred.  Please try again later."
+
+
+# Create the Gradio interface
+iface = gr.ChatInterface(
+    fn=chat_with_friendli,
+    title="Friendli AI Chatbot",
+    chatbot=gr.Chatbot(height=500),  # Adjust height as needed
+    textbox=gr.Textbox(placeholder="Type your message here...", container=False, scale=7),
+    theme="soft",
+    examples=["Explain the basics of cognitive behavioral therapy.", "What are some common psychological defense mechanisms?", "How can I deal with anxiety?"],
+    cache_examples=False, #remove the warnning from gradio
+    clear_btn="Clear",
+    retry_btn="Retry",
+    undo_btn="Undo",
+    submit_btn="Submit"
+)
+
+# Launch the Gradio app
 if __name__ == "__main__":
-    demo.launch()
+    iface.launch(debug=True)  # debug=True for development; set to False in production