invincible-jha commited on
Commit
d6fdb88
·
1 Parent(s): d28913a

Initialize AI models and agents with CPU-optimized models

Browse files
Files changed (2) hide show
  1. config/config.py +5 -5
  2. interface/app.py +78 -3
config/config.py CHANGED
@@ -6,22 +6,22 @@ def load_config() -> Dict:
6
  return {
7
  "MODEL_CONFIGS": {
8
  "conversation": {
9
- "model_id": "TheBloke/Llama-2-7B-Chat-GGUF",
10
  "max_length": 2048,
11
  "temperature": 0.7,
12
  },
13
  "emotion_detection": {
14
- "model_id": "bhadresh-savani/bert-base-uncased-emotion",
15
  },
16
  "speech_to_text": {
17
- "model_id": "openai/whisper-base",
18
  },
19
  "vision": {
20
- "model_id": "microsoft/resnet-50",
21
  }
22
  },
23
  "INTERFACE_CONFIG": {
24
- "theme": "light",
25
  "supported_languages": ["en"],
26
  "max_file_size_mb": 10,
27
  "supported_file_types": [
 
6
  return {
7
  "MODEL_CONFIGS": {
8
  "conversation": {
9
+ "model_id": "facebook/opt-125m",
10
  "max_length": 2048,
11
  "temperature": 0.7,
12
  },
13
  "emotion_detection": {
14
+ "model_id": "j-hartmann/emotion-english-distilroberta-base",
15
  },
16
  "speech_to_text": {
17
+ "model_id": "openai/whisper-tiny",
18
  },
19
  "vision": {
20
+ "model_id": "microsoft/resnet-18",
21
  }
22
  },
23
  "INTERFACE_CONFIG": {
24
+ "theme": "soft",
25
  "supported_languages": ["en"],
26
  "max_file_size_mb": 10,
27
  "supported_file_types": [
interface/app.py CHANGED
@@ -2,8 +2,13 @@ import os
2
  import torch
3
  import gradio as gr
4
  import logging
 
5
  from utils.log_manager import LogManager
6
  from utils.analytics_logger import AnalyticsLogger
 
 
 
 
7
 
8
  # Force CPU-only mode
9
  torch.cuda.is_available = lambda: False
@@ -21,9 +26,65 @@ class WellnessInterface:
21
  self.device = "cpu"
22
  self.logger.info("Using CPU-only mode")
23
 
 
 
 
 
 
 
24
  # Initialize interface
25
  self.setup_interface()
26
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
27
  def setup_interface(self):
28
  """Setup the Gradio interface components"""
29
  self.logger.info("Setting up interface components")
@@ -171,13 +232,27 @@ class WellnessInterface:
171
  }}
172
  )
173
 
174
- # Process the input (placeholder)
175
- response = "I understand and I'm here to help. Could you tell me more about how you're feeling?"
 
 
 
 
 
 
 
 
 
 
 
 
 
 
176
 
177
  # Add to chat history using message format
178
  history = history or []
179
  history.append({"role": "user", "content": text if text else "Sent media"})
180
- history.append({"role": "assistant", "content": response})
181
 
182
  return history, "" # Return empty string to clear text input
183
 
 
2
  import torch
3
  import gradio as gr
4
  import logging
5
+ from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
6
  from utils.log_manager import LogManager
7
  from utils.analytics_logger import AnalyticsLogger
8
+ from agents.conversation_agent import ConversationAgent
9
+ from agents.assessment_agent import AssessmentAgent
10
+ from agents.mindfulness_agent import MindfulnessAgent
11
+ from agents.crisis_agent import CrisisAgent
12
 
13
  # Force CPU-only mode
14
  torch.cuda.is_available = lambda: False
 
26
  self.device = "cpu"
27
  self.logger.info("Using CPU-only mode")
28
 
29
+ # Initialize models
30
+ self.initialize_models()
31
+
32
+ # Initialize agents
33
+ self.initialize_agents()
34
+
35
  # Initialize interface
36
  self.setup_interface()
37
 
38
+ def initialize_models(self):
39
+ """Initialize AI models"""
40
+ self.logger.info("Initializing AI models")
41
+ try:
42
+ # Initialize emotion detection model
43
+ self.emotion_model = pipeline(
44
+ "text-classification",
45
+ model=self.config["MODEL_CONFIGS"]["emotion_detection"]["model_id"],
46
+ device=self.device
47
+ )
48
+
49
+ # Initialize conversation model
50
+ self.conversation_tokenizer = AutoTokenizer.from_pretrained(
51
+ self.config["MODEL_CONFIGS"]["conversation"]["model_id"]
52
+ )
53
+ self.conversation_model = AutoModelForCausalLM.from_pretrained(
54
+ self.config["MODEL_CONFIGS"]["conversation"]["model_id"],
55
+ device_map={"": self.device}
56
+ )
57
+
58
+ self.logger.info("AI models initialized successfully")
59
+
60
+ except Exception as e:
61
+ self.logger.error(f"Error initializing models: {str(e)}")
62
+ raise
63
+
64
+ def initialize_agents(self):
65
+ """Initialize AI agents"""
66
+ self.logger.info("Initializing AI agents")
67
+ try:
68
+ # Initialize all agents
69
+ self.conversation_agent = ConversationAgent(
70
+ model_config=self.config["MODEL_CONFIGS"]
71
+ )
72
+ self.assessment_agent = AssessmentAgent(
73
+ model_config=self.config["MODEL_CONFIGS"]
74
+ )
75
+ self.mindfulness_agent = MindfulnessAgent(
76
+ model_config=self.config["MODEL_CONFIGS"]
77
+ )
78
+ self.crisis_agent = CrisisAgent(
79
+ model_config=self.config["MODEL_CONFIGS"]
80
+ )
81
+
82
+ self.logger.info("AI agents initialized successfully")
83
+
84
+ except Exception as e:
85
+ self.logger.error(f"Error initializing agents: {str(e)}")
86
+ raise
87
+
88
  def setup_interface(self):
89
  """Setup the Gradio interface components"""
90
  self.logger.info("Setting up interface components")
 
232
  }}
233
  )
234
 
235
+ # Analyze emotion if text is present
236
+ emotion = None
237
+ if text:
238
+ emotion_result = self.emotion_model(text)
239
+ emotion = emotion_result[0] if emotion_result else None
240
+ self.logger.info(f"Detected emotion: {emotion}")
241
+
242
+ # Route to appropriate agent based on content and emotion
243
+ if emotion and emotion.get("label") in ["anxiety", "fear", "panic"]:
244
+ response = self.crisis_agent.process_message(text)
245
+ elif "meditate" in text.lower() or "mindfulness" in text.lower():
246
+ response = self.mindfulness_agent.process_message(text)
247
+ elif "assess" in text.lower() or "check" in text.lower():
248
+ response = self.assessment_agent.process_message(text)
249
+ else:
250
+ response = self.conversation_agent.process_message(text)
251
 
252
  # Add to chat history using message format
253
  history = history or []
254
  history.append({"role": "user", "content": text if text else "Sent media"})
255
+ history.append({"role": "assistant", "content": response["message"]})
256
 
257
  return history, "" # Return empty string to clear text input
258