Update app.py

app.py

@@ -1,22 +1,23 @@
+!pip install streamlit transformers
 import streamlit as st
-from chatterbot import ChatBot
-from chatterbot.trainers import ChatterBotCorpusTrainer
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
-# Create a chatbot
-chatbot = ChatBot("KviGPT")
-trainer = ChatterBotCorpusTrainer(chatbot)
+device = "cuda"
+model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-v0.1")
+tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-v0.1")
 
-# Use the English corpus to train the chatbot
-trainer.train("chatterbot.corpus.english")
+model.to(device)
 
-# Define the function to get the chatbot response
-def get_response(input_text):
-    response = chatbot.get_response(input_text)
-    return str(response)
+def generate_text(prompt, max_new_tokens=100, do_sample=True):
+    model_inputs = tokenizer([prompt], return_tensors="pt").to(device)
+    generated_ids = model.generate(**model_inputs, max_new_tokens=max_new_tokens, do_sample=do_sample)
+    return tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
+
+st.title("KviGPT - Hugging Face Chat")
+
+user_input = st.text_input("You:", value="My favourite condiment is ")
 
-# Create the Streamlit app
-st.title("KviGPT Chatbot")
-user_input = st.text_area("You:", height=200)
 if st.button("Send"):
-    bot_response = get_response(user_input)
-    st.text_area("KviGPT:", value=bot_response, height=200)
+    prompt = user_input
+    model_response = generate_text(prompt)[0]
+    st.write("KviGPT:", model_response)