deploy at 2024-08-16 15:13:36.558496

main.py

@@ -1,17 +1,17 @@
 from fasthtml.common import *
-from fasthtml_hf import setup_hf_backup  # Importing setup_hf_backup
-import openai
+from fasthtml_hf import setup_hf_backup
 import uvicorn
+from transformers import pipeline
 
 # Set up the app, including daisyui and tailwind for the chat component
 tlink = Script(src="https://cdn.tailwindcss.com")
 dlink = Link(rel="stylesheet", href="https://cdn.jsdelivr.net/npm/[email protected]/dist/full.min.css")
 app = FastHTML(hdrs=(tlink, dlink, picolink))
 
-# Set up OpenAI API (https://openai.com/)
-# openai.api_key = os.getenv("OPENAI_API_KEY")
-messages = []
+# Set up the text generation pipeline
+pipe = pipeline("text-generation", model="ReliableAI/UCCIX-Llama2-13B-Instruct")
 
+messages = []
 
 # Chat message component (renders a chat bubble)
 def ChatMessage(msg):
@@ -21,12 +21,12 @@ def ChatMessage(msg):
                Div(msg['content'], cls=f"chat-bubble {bubble_class}"),
                cls=f"chat {chat_class}")
 
-# The input field for the user message. Also used to clear the 
+# The input field for the user message. Also used to clear the
 # input field after sending a message via an OOB swap
 def ChatInput():
-    return Input(type="text", name='msg', id='msg-input', 
-                 placeholder="Type a message", 
-                 cls="input input-bordered w-full", hx_swap_oob='true')
+    return Input(type="text", name='msg', id='msg-input',
+                  placeholder="Type a message",
+                  cls="input input-bordered w-full", hx_swap_oob='true')
 
 # The main screen
 @app.route("/")
@@ -45,24 +45,20 @@ def get():
 def post(msg:str):
     messages.append({"role":"user", "content":msg})
     
-    # Get response from OpenAI API
-    response = openai.chat.completions.create(
-        model="gpt-3.5-turbo",
-        messages=[
-            {"role": "system", "content": "You are a helpful and concise assistant."},
-            *messages
-        ]
-    )
+    # Get response from the text generation pipeline
+    full_prompt = "You are a helpful and concise assistant.\n\n"
+    for m in messages:
+        full_prompt += f"{m['role'].capitalize()}: {m['content']}\nAssistant: "
     
-    assistant_msg = response.choices[0].message.content.strip()
+    response = pipe(full_prompt, max_length=2048, num_return_sequences=1)
+    
+    assistant_msg = response[0]['generated_text'].split("Assistant: ")[-1].strip()
     messages.append({"role":"assistant", "content":assistant_msg})
     
     return (ChatMessage(messages[-2]), # The user's message
             ChatMessage(messages[-1]), # The chatbot's response
             ChatInput()) # And clear the input field via an OOB swap
 
-
-#serve()
 if __name__ == "__main__":
     setup_hf_backup(app)
     uvicorn.run(app, host="0.0.0.0", port=7860)