Update app.py

app.py

@@ -1,37 +1,119 @@
 import gradio as gr
-from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import pipeline
+import torch
+from fastapi import FastAPI, Request
+from fastapi.middleware.cors import CORSMiddleware
+import json
+
+# Initialize FastAPI
+app = FastAPI()
+
+# Add CORS middleware to allow requests from your React Native app
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],  # In production, restrict this to your app's domain
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
 
 # Load chatbot model
+print("Loading chatbot model...")
 chatbot_model = "microsoft/DialoGPT-medium"
 tokenizer = AutoTokenizer.from_pretrained(chatbot_model)
 model = AutoModelForCausalLM.from_pretrained(chatbot_model)
 
 # Load emotion detection model
-emotion_pipeline = pipeline("text-classification", model="bhadresh-savani/distilbert-base-uncased-emotion")
+print("Loading emotion detection model...")
+emotion_pipeline = pipeline(
+    "text-classification", 
+    model="bhadresh-savani/distilbert-base-uncased-emotion",
+    return_all_scores=False
+)
+
+# Store conversation history
+chat_history = {}
+
+# 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")
+    
+    # Process input and get response
+    response, emotion, score = generate_response(user_input, session_id)
+    
+    return {
+        "response": response,
+        "emotion": emotion,
+        "score": float(score),
+        "session_id": session_id
+    }
 
-def generate_response(user_input):
-    # 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)
+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] = []
+    
+    # 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
-    emotion_result = emotion_pipeline(user_input)
-    emotion = emotion_result[0]["label"]
-    score = emotion_result[0]["score"]
+    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}"
+
+# Create Gradio interface
 iface = gr.Interface(
-    fn=generate_response,
-    inputs=gr.Textbox(label="Enter your message"),
+    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")
     ],
-    live=False
+    title="Mental Health Chatbot",
+    description="A simple mental health chatbot with emotion detection",
+    allow_flagging="never"
 )
 
-iface.launch()
+# Mount the Gradio app to FastAPI
+app = gr.mount_gradio_app(app, iface, path="/")
+
+# Run the app
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)