Arjun Moorthy commited on
Commit
6ab8dba
Β·
1 Parent(s): 2af3862

Simplify to use only BioMistral base model, remove all adapter references

Browse files
Files changed (1) hide show
  1. app.py +12 -12
app.py CHANGED
@@ -19,11 +19,11 @@ def force_gpu_detection():
19
 
20
  class HFMedicalChatbot:
21
  def __init__(self):
22
- # BioMistral model configuration
23
- self.model_id = "BioMistral/BioMistral-7B" # Direct BioMistral model
24
 
25
  print("πŸ”„ Initializing HF Medical Chatbot")
26
- print(f"πŸ“¦ Loading model: {self.model_id}")
27
 
28
  # Force GPU detection first
29
  try:
@@ -33,12 +33,12 @@ class HFMedicalChatbot:
33
  print(f"⚠️ GPU detection error: {e}")
34
  gpu_available = torch.cuda.is_available()
35
 
36
- self._load_model(gpu_available)
37
 
38
- def _load_model(self, gpu_available):
39
- """Load the BioMistral medical model directly"""
40
  try:
41
- print("πŸ”„ Loading BioMistral model...")
42
 
43
  # Determine device strategy
44
  if gpu_available and torch.cuda.is_available():
@@ -51,16 +51,16 @@ class HFMedicalChatbot:
51
  print("πŸ’» Loading BioMistral model on CPU...")
52
 
53
  # Load tokenizer
54
- print(f"πŸ“ Loading tokenizer: {self.model_id}")
55
  self.tokenizer = AutoTokenizer.from_pretrained(
56
- self.model_id,
57
  trust_remote_code=True
58
  )
59
 
60
  # Load the model
61
- print(f"πŸ“¦ Loading model: {self.model_id}")
62
  self.model = MistralForCausalLM.from_pretrained(
63
- self.model_id,
64
  trust_remote_code=True,
65
  device_map="auto",
66
  torch_dtype=dtype,
@@ -71,7 +71,7 @@ class HFMedicalChatbot:
71
  if self.tokenizer.pad_token is None:
72
  self.tokenizer.pad_token = self.tokenizer.eos_token
73
 
74
- print(f"βœ… BioMistral model loaded successfully on {device.upper()}!")
75
 
76
  except Exception as e:
77
  print(f"❌ Error loading BioMistral model: {e}")
 
19
 
20
  class HFMedicalChatbot:
21
  def __init__(self):
22
+ # BioMistral base model configuration
23
+ BASE = "BioMistral/BioMistral-7B"
24
 
25
  print("πŸ”„ Initializing HF Medical Chatbot")
26
+ print(f"πŸ“¦ Loading base model: {BASE}")
27
 
28
  # Force GPU detection first
29
  try:
 
33
  print(f"⚠️ GPU detection error: {e}")
34
  gpu_available = torch.cuda.is_available()
35
 
36
+ self._load_model(BASE, gpu_available)
37
 
38
+ def _load_model(self, model_id, gpu_available):
39
+ """Load the BioMistral base model"""
40
  try:
41
+ print("πŸ”„ Loading BioMistral base model...")
42
 
43
  # Determine device strategy
44
  if gpu_available and torch.cuda.is_available():
 
51
  print("πŸ’» Loading BioMistral model on CPU...")
52
 
53
  # Load tokenizer
54
+ print(f"πŸ“ Loading tokenizer: {model_id}")
55
  self.tokenizer = AutoTokenizer.from_pretrained(
56
+ model_id,
57
  trust_remote_code=True
58
  )
59
 
60
  # Load the model
61
+ print(f"πŸ“¦ Loading model: {model_id}")
62
  self.model = MistralForCausalLM.from_pretrained(
63
+ model_id,
64
  trust_remote_code=True,
65
  device_map="auto",
66
  torch_dtype=dtype,
 
71
  if self.tokenizer.pad_token is None:
72
  self.tokenizer.pad_token = self.tokenizer.eos_token
73
 
74
+ print(f"βœ… BioMistral base model loaded successfully on {device.upper()}!")
75
 
76
  except Exception as e:
77
  print(f"❌ Error loading BioMistral model: {e}")