Update app.py

app.py

@@ -1,15 +1,15 @@
-import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer
-from transformers import pipeline
-import torch
 from fastapi import FastAPI, Request
 from fastapi.middleware.cors import CORSMiddleware
-import json
+from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
+from pydantic import BaseModel
+import uuid
 
-# Initialize FastAPI
-app = FastAPI()
+# Initialize FastAPI app
+app = FastAPI(title="Mental Health Chatbot API", 
+              description="API for mental health chatbot with emotion detection",
+              version="1.0.0")
 
-# Add CORS middleware to allow requests from your React Native app
+# Add CORS middleware
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],  # In production, restrict this to your app's domain
@@ -19,101 +19,82 @@ app.add_middleware(
 )
 
 # Load chatbot model
-print("Loading chatbot model...")
+print("Loading DialoGPT model...")
 chatbot_model = "microsoft/DialoGPT-medium"
 tokenizer = AutoTokenizer.from_pretrained(chatbot_model)
 model = AutoModelForCausalLM.from_pretrained(chatbot_model)
 
 # Load emotion detection model
 print("Loading emotion detection model...")
-emotion_pipeline = pipeline(
-    "text-classification", 
-    model="bhadresh-savani/distilbert-base-uncased-emotion",
-    return_all_scores=False
-)
+emotion_pipeline = pipeline("text-classification", model="j-hartmann/emotion-english-distilroberta-base")
+
+# Store chat histories by session ID
+chat_histories = {}
+
+# Request models
+class ChatRequest(BaseModel):
+    message: str
+    session_id: str = None
 
-# Store conversation history
-chat_history = {}
+class ChatResponse(BaseModel):
+    response: str
+    emotion: str
+    emotion_score: float
+    session_id: str
 
-# Define API endpoint for chatbot
-@app.post("/api/chat")
-async def chat_endpoint(request: Request):
-    data = await request.json()
-    user_input = data.get("message", "")
-    session_id = data.get("session_id", "default")
+@app.get("/")
+async def root():
+    return {"message": "Mental Health Chatbot API is running. Use /api/chat to interact with the chatbot."}
+
+@app.post("/api/chat", response_model=ChatResponse)
+async def chat(request: ChatRequest):
+    # Create a new session ID if not provided
+    session_id = request.session_id if request.session_id else str(uuid.uuid4())
+    
+    # Initialize chat history for new sessions
+    if session_id not in chat_histories:
+        chat_histories[session_id] = []
     
-    # Process input and get response
-    response, emotion, score = generate_response(user_input, session_id)
+    user_input = request.message
     
+    # Generate chatbot response
+    input_ids = tokenizer.encode(user_input + tokenizer.eos_token, return_tensors="pt")
+    output = model.generate(input_ids, max_length=200, pad_token_id=tokenizer.eos_token_id)
+    response = tokenizer.decode(output[:, input_ids.shape[-1]:][0], skip_special_tokens=True)
+    
+    # Detect emotion
+    emotion_result = emotion_pipeline(user_input)[0]
+    emotion = emotion_result["label"]
+    score = emotion_result["score"]
+    
+    # Store chat history
+    chat_histories[session_id].append({"role": "user", "content": user_input})
+    chat_histories[session_id].append({"role": "bot", "content": response})
+    
+    # Return the response
     return {
         "response": response,
         "emotion": emotion,
-        "score": float(score),
+        "emotion_score": float(score),
         "session_id": session_id
     }
 
-def generate_response(user_input, session_id="default"):
-    # Initialize chat history for new sessions
-    if session_id not in chat_history:
-        chat_history[session_id] = []
+@app.get("/api/history/{session_id}")
+async def get_chat_history(session_id: str):
+    if session_id not in chat_histories:
+        return {"error": "Session not found"}
     
-    # Format the input with chat history
-    bot_input_ids = tokenizer.encode(user_input + tokenizer.eos_token, return_tensors="pt")
-    
-    # Append user input to chat history
-    chat_history[session_id].append(bot_input_ids)
-    
-    # Generate a response
-    with torch.no_grad():
-        chat_history_ids = model.generate(
-            bot_input_ids,
-            max_length=200,
-            pad_token_id=tokenizer.eos_token_id,
-            no_repeat_ngram_size=3,
-            do_sample=True,
-            top_k=50,
-            top_p=0.95,
-            temperature=0.7
-        )
-    
-    # Decode the response
-    response = tokenizer.decode(chat_history_ids[:, bot_input_ids.shape[-1]:][0], skip_special_tokens=True)
-    
-    # Detect emotion
-    try:
-        emotion_result = emotion_pipeline(user_input)[0]
-        emotion = emotion_result["label"]
-        score = emotion_result["score"]
-    except Exception as e:
-        print(f"Error detecting emotion: {e}")
-        emotion = "unknown"
-        score = 0.0
-    
-    return response, emotion, score
-
-# Gradio Interface
-def gradio_generate(user_input):
-    response, emotion, score = generate_response(user_input)
-    return response, emotion, f"{score:.4f}"
+    return {"session_id": session_id, "history": chat_histories[session_id]}
 
-# Create Gradio interface
-iface = gr.Interface(
-    fn=gradio_generate,
-    inputs=gr.Textbox(label="Enter your message", placeholder="Type your message here..."),
-    outputs=[
-        gr.Textbox(label="Chatbot Response"),
-        gr.Textbox(label="Emotion Detected"),
-        gr.Textbox(label="Emotion Score")
-    ],
-    title="Mental Health Chatbot",
-    description="A simple mental health chatbot with emotion detection",
-    allow_flagging="never"
-)
-
-# Mount the Gradio app to FastAPI
-app = gr.mount_gradio_app(app, iface, path="/")
+@app.delete("/api/history/{session_id}")
+async def clear_chat_history(session_id: str):
+    if session_id in chat_histories:
+        chat_histories.pop(session_id)
+        return {"message": f"Chat history for session {session_id} cleared"}
+    
+    return {"error": "Session not found"}
 
-# Run the app
+# Launch the API server
 if __name__ == "__main__":
     import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=7860)