Update app.py

app.py

@@ -1,27 +1,20 @@
 import gradio as gr
-from typing import Optional, Tuple
-from langchain.llm import LLMChain
-from langchain.memory import ConversationBufferWindowMemory
-from langchain.huggingface_hub import HuggingFaceHub
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
-# Set up the language model chain
-prompt = "Instructions: You are SplitticAI. You answer questions exactly like people ask them. You were made by SplitticHost. You impersonate yourself as an AI chatbot.\n\n"
-llm = HuggingFaceHub(repo_id="google/flan-t5-xxl", model_kwargs={"temperature":1e-10})
-llm_chain = LLMChain(
-    llm=llm,
-    prompt=prompt,
-    verbose=True,
-    memory=ConversationBufferWindowMemory(k=2)
-)
+# Load the model and tokenizer
+model_name = "google/flax-t5-xxl-qa-121k"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name)
 
 # Define the chat function
-def chat(
-    inp: str, history: Optional[Tuple[str, str]], chain: Optional[ConversationChain]
-):
-    history = history or []
-    output = llm_chain.predict(human_input=inp)
-    history.append((inp, output))
-    return history, history
+def chat(message):
+    # Encode the user's message
+    inputs = tokenizer.encode(message, return_tensors="pt")
+    # Generate a response from the model
+    outputs = model.generate(inputs, max_length=1024, pad_token_id=tokenizer.eos_token_id)
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    # Return the response
+    return response
 
 # Set up the Gradio interface
 block = gr.Blocks(css=".gradio-container {background-color: lightgray}")
@@ -52,13 +45,12 @@ with block:
     gr.HTML("Ask SplitticAI anything and get an answer!")
 
     gr.HTML(
-        "<center>Powered by SplitticHost</center>"
+        "<center>Powered by <a href='https://huggingface.co/google/flax-t5-xxl-qa-121k'>google/flax-t5-xxl-qa-121k</a></center>"
     )
 
     state = gr.State()
     agent_state = gr.State()
 
-    submit.click(chat, inputs=[message, state, agent_state], outputs=[chatbot, state])
-    message.submit(chat, inputs=[message, state, agent_state], outputs=[chatbot, state])
+    submit.click(chat, inputs=[message], outputs=[chatbot])
 
 block.launch(debug=True)