Update app.py

app.py

@@ -1,25 +1,29 @@
-from llama_cpp import Llama
+from transformers import GPT2LMHeadModel, GPT2Tokenizer
 import gradio as gr
 
-# Load GGUF model
-llm = Llama(
-    model_path="Mental-Health-Chatbot.i1-IQ1_M.gguf", # Make sure this filename matches exactly
-    n_ctx=2048,
-    n_threads=4,
-)
+# Load the tokenizer and model from Hugging Face
+tokenizer = GPT2Tokenizer.from_pretrained("distilgpt2")
+model = GPT2LMHeadModel.from_pretrained("distilgpt2")
+
+# Ensure the model doesn't generate any special tokens like <pad>
+tokenizer.pad_token = tokenizer.eos_token
 
-# Chat logic
 def chat(message, history):
+    # Prepare the conversation history
     full_prompt = ""
     for user, bot in history:
         full_prompt += f"User: {user}\nBot: {bot}\n"
     full_prompt += f"User: {message}\nBot:"
 
-    output = llm(full_prompt, max_tokens=128, stop=["User:", "\n"], echo=False)
-    reply = output["choices"][0]["text"].strip()
+    # Tokenize the input and generate a response
+    inputs = tokenizer(full_prompt, return_tensors="pt")
+    outputs = model.generate(inputs["input_ids"], max_length=150, num_return_sequences=1, no_repeat_ngram_size=2)
+    reply = tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+    # Extract only the new reply
+    reply = reply.split("Bot:")[-1].strip()
+
     return reply
 
-# Chat Interface
-gr.ChatInterface(fn=chat, title="Mental Health Llama Chatbot").launch(
-    server_name="0.0.0.0", server_port=7860
-)
+# Set up the Gradio interface
+gr.ChatInterface(fn=chat, title="Simple Chatbot with DistilGPT-2").launch()