Update app.py

app.py

@@ -1,38 +1,40 @@
 import os
+import gradio as gr
 from langchain_ai21 import ChatAI21
-from langchain.prompts import PromptTemplate
-from langchain.chains import LLMChain
-
-# 1. Set up your AI21 API key
-os.environ["AI21_API_KEY"] = "your-ai21-api-key"
-
-# 2. Create a prompt template for the chatbot
-prompt = PromptTemplate(
-    input_variables=["user_input"],
-    template="""
-    You are a helpful and friendly chatbot. Respond concisely and informatively.
-
-    User: {user_input}
-    Chatbot:
-    """
-)
-llm = ChatAI21(model="jamba-instruct", temperature=0)
-
-from langchain.memory import ConversationBufferMemory
+from ai21 import AI21Client
+from ai21.models.chat import ChatMessage, DocumentSchema
 
-memory = ConversationBufferMemory()
-chat_chain = LLMChain(llm=llm, prompt=prompt, memory=memory)
-
-llm.streaming = True
-
-import gradio as gr
+# Set your AI21 API key
+os.environ["AI21_API_KEY"] = "8T6NvXgGjhtlh9bh65jsNqb584BOorNM"
+client = AI21Client(api_key="8T6NvXgGjhtlh9bh65jsNqb584BOorNM")
+# Initialize the Jamba model
+chatbot = ChatAI21(model="jamba-instruct", temperature=0.7)
 
+# Define the function to handle chat
 def chatbot_response(user_input):
-    return chat_chain.invoke({"user_input": user_input})
-
-gr.Interface(
+    # Wrap input into a dictionary with the expected format
+    messages = [ChatMessage(role='system', content='You are a concise factual based question answering assistant.'),
+            ChatMessage(role='user', content=user_input)
+            ]
+    response = client.chat.completions.create(messages=messages,
+                                          model='jamba-1.5-large',
+                                          # max_tokens=4096,
+                                          # temperature=0.4,
+                                          # top_p=1.0,
+                                          # stop = [], ## ['####', '\n'],
+                                          # n=1,
+                                          # stream = False
+                                          )
+    return response.choices[0].message.content
+# Create the Gradio interface
+interface = gr.Interface(
     fn=chatbot_response,
     inputs="text",
     outputs="text",
-    title="AI Chatbot with Jamba"
-).launch()
+    title="Jamba Chatbot",
+    description="A simple chatbot using AI21 Labs' Jamba technology."
+)
+
+# Launch the Gradio app
+if __name__ == "__main__":
+    interface.launch()