Spaces:
Sleeping
Sleeping
File size: 6,362 Bytes
37244c4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 |
"""
API Server for MedLLaMA2 Medical Chatbot
This file provides REST API endpoints that can be used by external applications
while the main app.py provides the Gradio interface.
"""
import os
import threading
from flask import Flask, request, jsonify, Response
from flask_cors import CORS
import json
import time
import re
# Import the model and functions from the main app
from app import load_model, generate_response, get_model_info
from config import GENERATION_DEFAULTS
# Initialize Flask app
app = Flask(__name__)
CORS(app) # Enable CORS for all routes
# Initialize model in a separate thread
def init_model():
print("🔄 Loading model in API server...")
load_model()
print("✅ Model loaded in API server")
# Start model loading
model_thread = threading.Thread(target=init_model)
model_thread.start()
@app.route('/health', methods=['GET'])
def health_check():
"""Health check endpoint"""
return jsonify({
'status': 'ok',
'model_loaded': get_model_info() != "No model loaded",
'model_info': get_model_info(),
'timestamp': time.time()
})
@app.route('/chat', methods=['POST'])
def chat_endpoint():
"""Main chat endpoint for medical questions"""
try:
data = request.get_json()
if not data or 'message' not in data:
return jsonify({'error': 'No message provided'}), 400
message = data['message'].strip()
if not message:
return jsonify({'error': 'Empty message'}), 400
# Get optional parameters
max_tokens = data.get('max_tokens', GENERATION_DEFAULTS['max_new_tokens'])
temperature = data.get('temperature', GENERATION_DEFAULTS['temperature'])
top_p = data.get('top_p', GENERATION_DEFAULTS['top_p'])
# Check for non-medical topics
non_medical_patterns = [
r'\b(java|javascript|python|c\+\+|c#|programming|coding|computer|software)\b',
r'\b(cook|recipe|food recipe|baking)\b',
r'\b(math problem|finance|stock market|weather|movie|book|travel)\b'
]
is_non_medical = any(re.search(pattern, message, re.IGNORECASE) for pattern in non_medical_patterns)
# Medical exceptions
medical_exceptions = [
r'medical (history|coding|program|software|algorithm)',
r'health (history|software|recipe)',
r'(food allergy|diet recipe|patient story|medical story)'
]
is_medical_exception = any(re.search(pattern, message, re.IGNORECASE) for pattern in medical_exceptions)
if is_non_medical and not is_medical_exception:
return jsonify({
'response': "I'm a medical assistant designed to provide health-related information. I'm not able to help with programming, cooking, or other non-medical topics. If you have any questions about health, medicine, symptoms, or wellness, I'd be happy to assist you! 😊",
'timestamp': time.time()
})
# Generate medical response
response = generate_response(
message,
max_tokens=int(max_tokens),
temperature=float(temperature),
top_p=float(top_p)
)
# Return the response
return jsonify({
'response': response,
'timestamp': time.time(),
'model_info': get_model_info()
})
except Exception as e:
print(f"Error in chat endpoint: {str(e)}")
return jsonify({
'error': 'Internal server error',
'details': str(e)
}), 500
@app.route('/stream', methods=['POST'])
def stream_chat():
"""Streaming chat endpoint"""
try:
data = request.get_json()
if not data or 'message' not in data:
return jsonify({'error': 'No message provided'}), 400
message = data['message'].strip()
if not message:
return jsonify({'error': 'Empty message'}), 400
def generate_stream():
try:
# Get parameters
max_tokens = data.get('max_tokens', GENERATION_DEFAULTS['max_new_tokens'])
temperature = data.get('temperature', GENERATION_DEFAULTS['temperature'])
top_p = data.get('top_p', GENERATION_DEFAULTS['top_p'])
# Generate response in chunks
response = generate_response(
message,
max_tokens=int(max_tokens),
temperature=float(temperature),
top_p=float(top_p)
)
# Stream the response word by word
words = response.split()
for i, word in enumerate(words):
chunk_data = {
'chunk': word + (' ' if i < len(words) - 1 else ''),
'status': 'streaming'
}
yield f"data: {json.dumps(chunk_data)}\n\n"
time.sleep(0.05) # Small delay for streaming effect
# Send completion signal
end_data = {
'complete': True,
'fullResponse': response
}
yield f"event: end\ndata: {json.dumps(end_data)}\n\n"
except Exception as e:
error_data = {
'error': 'Stream error',
'details': str(e)
}
yield f"event: error\ndata: {json.dumps(error_data)}\n\n"
return Response(
generate_stream(),
content_type='text/event-stream',
headers={
'Cache-Control': 'no-cache',
'Connection': 'keep-alive',
'Access-Control-Allow-Origin': '*',
'Access-Control-Allow-Headers': 'Content-Type, Authorization'
}
)
except Exception as e:
return jsonify({'error': str(e)}), 500
if __name__ == "__main__":
# For local development
port = int(os.environ.get("API_PORT", 8000))
print(f"🚀 Starting API server on port {port}")
app.run(host="0.0.0.0", port=port, debug=False)
|