Update app.py

app.py

@@ -1,13 +1,15 @@
 import gradio as gr
-from transformers import pipeline
+from transformers import AutoTokenizer, AutoModelForCausalLM
 from datetime import datetime, timezone
 
-# Initialize the model
-generator = pipeline('text-generation', model='facebook/opt-350m')  # We'll update this with your medical model later
+# Initialize the model and tokenizer
+model_id = "CohereLabs/c4ai-command-a-03-2025"
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(model_id)
 
 def get_timestamp():
     """Get current UTC datetime in specified format"""
-    return datetime.now(timezone.utc).strftime('%Y-%m-%d %H:%M:%S')
+    return datetime.now(timezone.utc).strftime('%Y-%m-%d %H:%M:%SS')
 
 def format_system_info():
     """Format system information header"""
@@ -23,31 +25,29 @@ def chat(message, history):
     # Add system information
     system_info = format_system_info()
     
-    # Format the prompt with system info
-    prompt = f"{system_info}\nPatient Query: {message}\nMedical AI Assistant:"
-
-    try:
-        # Generate response
-        response = generator(
-            prompt,
-            max_length=512,
-            temperature=0.7,
-            top_p=0.95,
-            do_sample=True
-        )[0]['generated_text']
-        
-        # Extract and format the response
-        ai_response = response[len(prompt):].strip()
-        formatted_response = f"{system_info}\n{ai_response}"
-        
-        history.append((message, formatted_response))
-        return history
-    except Exception as e:
-        return [(message, f"Error: {str(e)}")]
+    # Format messages with the c4ai-command-a-03-2025 chat template
+    messages = [{"role": "user", "content": message}]
+    input_ids = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True)
+    
+    # Generate response
+    gen_tokens = model.generate(
+        input_ids,
+        max_new_tokens=100,
+        do_sample=True,
+        temperature=0.3,
+    )
+    
+    # Decode response
+    gen_text = tokenizer.decode(gen_tokens[0])
+    
+    # Format the full response with system info
+    formatted_response = f"{system_info}\n{gen_text}"
+    history.append((message, formatted_response))
+    return history
 
 # Create custom theme
 theme = gr.themes.Default().set(
-    body_background_fill="#f0f8ff",  # Light blue background
+    body_background_fill="#f0f8ff",
     block_background_fill="#ffffff",
     block_border_width="1px",
     block_border_color="#2c3e50",