Update app.py

app.py

@@ -1,41 +1,91 @@
+Hugging Face's logo
+Hugging Face
+Search models, datasets, users...
+Models
+Datasets
+Spaces
+Docs
+Solutions
+Pricing
+
+
+
+Spaces:
+
+CaioMartins1
+/
+test
+
+
+like
+0
+
+Logs
+App
+Files
+Community
+Settings
+test
+/
+app.py
+CaioMartins1's picture
+CaioMartins1
+Update app.py
+604235c
+about 3 hours ago
+raw
+history
+blame
+edit
+delete
+No virus
+1.59 kB
 import torch
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
+from transformers import pipeline, BertTokenizer, BertForQuestionAnswering
 from datasets import load_dataset
 
-# Load the depression advice dataset
-dataset = load_dataset("ziq/depression_advice")
+# Load the dataset
+advice_dataset = load_dataset("ziq/depression_advice")
+
+# 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)
 
-# Load the fine-tuned BERT model
-model_name = "mrm8488/distilroberta-base-finetuned-suicide-depression"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForSequenceClassification.from_pretrained(model_name)
+# Extract context and messages
+contexts = advice_dataset["train"]["text"]
 
-# Function to generate depression advice using the fine-tuned BERT model
-def generate_advice(prompt):
-    # Tokenize the input prompt
-    inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True)
+# Define a function to generate answers
+def generate_answer(messages):
+    # If messages is a list, use the first message
+    if isinstance(messages, list):
+        messages = messages[0]
+    
+    # Tokenize the input message
+    inputs = tokenizer(messages, return_tensors="pt")
 
-    # Get model predictions
-    outputs = model(**inputs)
-    logits = outputs.logits
+    # Use the fine-tuned BERT model to generate the answer for the single message
+    with torch.no_grad():
+        outputs = model(**inputs)
 
-    # Get the predicted label
-    predicted_label_id = torch.argmax(logits, dim=1).item()
+    # 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 the generated advice based on the predicted label
-    advice = dataset["train"][predicted_label_id]["text"]
-    return advice
+    return answer if answer else "No answer found."
 
-# Create the Gradio interface
+# Create a Gradio interface
 iface = gr.Interface(
-    fn=generate_advice,
-    inputs=gr.Textbox(),
-    outputs=gr.Textbox(),
-    live=True,
+    fn=generate_answer,
+    inputs=[
+        gr.Textbox(type="text", label="Message"),
+    ],
+    outputs=gr.Textbox(type="text", label="Answer"),
     title="Depression Advice Generator",
     description="Enter your feelings, and get supportive advice generated by a fine-tuned BERT model.",
 )
 
-# Launch the Gradio interface
+# Launch the interface
 iface.launch()