Update app.py

app.py

@@ -1,18 +1,17 @@
 import gradio as gr
-from transformers import pipeline
+from transformers import pipeline, BertTokenizer, BertForQuestionAnswering
 from datasets import load_dataset
 
 # Load the dataset
 advice_dataset = load_dataset("ziq/depression_advice")
-depression_dataset = load_dataset("ShreyaR/DepressionDetection")
 
-# Load the model
-model_name = "mrm8488/distilroberta-base-finetuned-suicide-depression"
-qa_pipeline = pipeline("question-answering", model=model_name)
+# Load the fine-tuned BERT model and tokenizer
+model_dir = "./bert-finetuned-depression"
+model = BertForQuestionAnswering.from_pretrained(model_dir)
+tokenizer = BertTokenizer.from_pretrained(model_dir)
 
 # Extract context and messages
 contexts = advice_dataset["train"]["text"]
-messages = depression_dataset["train"]["clean_text"]
 
 # Define a function to generate answers
 def generate_answer(messages):
@@ -20,11 +19,19 @@ def generate_answer(messages):
     if isinstance(messages, list):
         messages = messages[0]
     
-    # Use the QA model to generate the answer for the single message
-    results = qa_pipeline(question=messages, context=contexts)
-    
-    # Return the answer
-    return results["answer"] if results["answer"] else "No answer found."
+    # Tokenize the input message
+    inputs = tokenizer(messages, return_tensors="pt")
+
+    # Use the fine-tuned BERT model to generate the answer for the single message
+    with torch.no_grad():
+        outputs = model(**inputs)
+
+    # Decode the output and return the answer
+    answer_start = torch.argmax(outputs.start_logits)
+    answer_end = torch.argmax(outputs.end_logits) + 1
+    answer = tokenizer.decode(inputs["input_ids"][0][answer_start:answer_end])
+
+    return answer if answer else "No answer found."
 
 # Create a Gradio interface
 iface = gr.Interface(