Deploy NEED AI API - app.py

app.py

@@ -1,23 +1,13 @@
 #!/usr/bin/env python3
-"""
-NEED AI - Production Flask API with Direct Model Loading
-"""
-
 from flask import Flask, request, jsonify
 from flask_cors import CORS
-from transformers import (
-    T5ForConditionalGeneration, 
-    T5Tokenizer,
-    AutoModelForSequenceClassification,
-    AutoTokenizer
-)
+from transformers import T5ForConditionalGeneration, T5Tokenizer, AutoModelForSequenceClassification, AutoTokenizer
 from sentence_transformers import SentenceTransformer
 import torch
 import torch.nn.functional as F
 from sklearn.metrics.pairwise import cosine_similarity
 import logging
 import os
-from functools import lru_cache
 import time
 
 logging.basicConfig(level=logging.INFO)
@@ -26,300 +16,130 @@ logger = logging.getLogger(__name__)
 app = Flask(__name__)
 CORS(app)
 
-HF_USERNAME = os.getenv("HF_USERNAME", "yogami9")
+HF_USERNAME = "yogami9"
 DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-logger.info(f"Using device: {DEVICE}")
 
 class ModelCache:
     def __init__(self):
         self.models = {}
         self.tokenizers = {}
-        logger.info("Model cache initialized")
-        
-    @lru_cache(maxsize=1)
+    
     def get_category_model(self):
         if 'category' not in self.models:
             logger.info("Loading Category model...")
-            model_name = f"{HF_USERNAME}/need-category-recommendation"
-            self.models['category'] = T5ForConditionalGeneration.from_pretrained(model_name).to(DEVICE)
-            self.tokenizers['category'] = T5Tokenizer.from_pretrained(model_name)
-            logger.info("✅ Category model loaded")
+            self.models['category'] = T5ForConditionalGeneration.from_pretrained(f"{HF_USERNAME}/need-category-recommendation").to(DEVICE)
+            self.tokenizers['category'] = T5Tokenizer.from_pretrained(f"{HF_USERNAME}/need-category-recommendation")
         return self.models['category'], self.tokenizers['category']
     
-    @lru_cache(maxsize=1)
     def get_chat_model(self):
         if 'chat' not in self.models:
             logger.info("Loading Chat model...")
-            model_name = f"{HF_USERNAME}/need-chat-support"
-            self.models['chat'] = T5ForConditionalGeneration.from_pretrained(model_name).to(DEVICE)
-            self.tokenizers['chat'] = T5Tokenizer.from_pretrained(model_name)
-            logger.info("✅ Chat model loaded")
+            self.models['chat'] = T5ForConditionalGeneration.from_pretrained(f"{HF_USERNAME}/need-chat-support").to(DEVICE)
+            self.tokenizers['chat'] = T5Tokenizer.from_pretrained(f"{HF_USERNAME}/need-chat-support")
         return self.models['chat'], self.tokenizers['chat']
     
-    @lru_cache(maxsize=1)
     def get_service_model(self):
         if 'service' not in self.models:
             logger.info("Loading Service model...")
-            model_name = f"{HF_USERNAME}/need-service-description"
-            self.models['service'] = T5ForConditionalGeneration.from_pretrained(model_name).to(DEVICE)
-            self.tokenizers['service'] = T5Tokenizer.from_pretrained(model_name)
-            logger.info("✅ Service model loaded")
+            self.models['service'] = T5ForConditionalGeneration.from_pretrained(f"{HF_USERNAME}/need-service-description").to(DEVICE)
+            self.tokenizers['service'] = T5Tokenizer.from_pretrained(f"{HF_USERNAME}/need-service-description")
         return self.models['service'], self.tokenizers['service']
     
-    @lru_cache(maxsize=1)
     def get_search_model(self):
         if 'search' not in self.models:
             logger.info("Loading Search model...")
-            model_name = f"{HF_USERNAME}/need-semantic-search"
-            self.models['search'] = SentenceTransformer(model_name)
-            logger.info("✅ Search model loaded")
+            self.models['search'] = SentenceTransformer(f"{HF_USERNAME}/need-semantic-search")
         return self.models['search']
     
-    @lru_cache(maxsize=1)
     def get_moderation_model(self):
         if 'moderation' not in self.models:
             logger.info("Loading Moderation model...")
-            model_name = f"{HF_USERNAME}/need-content-moderation"
-            self.models['moderation'] = AutoModelForSequenceClassification.from_pretrained(model_name).to(DEVICE)
-            self.tokenizers['moderation'] = AutoTokenizer.from_pretrained(model_name)
-            logger.info("✅ Moderation model loaded")
+            self.models['moderation'] = AutoModelForSequenceClassification.from_pretrained(f"{HF_USERNAME}/need-content-moderation").to(DEVICE)
+            self.tokenizers['moderation'] = AutoTokenizer.from_pretrained(f"{HF_USERNAME}/need-content-moderation")
         return self.models['moderation'], self.tokenizers['moderation']
 
-model_cache = ModelCache()
+cache = ModelCache()
 
-@app.route('/', methods=['GET'])
+@app.route('/')
 def home():
-    return jsonify({
-        'name': 'NEED AI API',
-        'version': '1.0.0',
-        'status': 'running',
-        'endpoints': {
-            'health': '/health',
-            'category': '/api/category',
-            'chat': '/api/chat',
-            'service': '/api/service',
-            'search': '/api/search',
-            'moderate': '/api/moderate',
-            'batch': '/api/batch'
-        },
-        'documentation': 'https://github.com/Need-Service-App/need-ai-model'
-    })
+    return jsonify({'name': 'NEED AI API', 'status': 'running', 'models': 5})
 
-@app.route('/health', methods=['GET'])
+@app.route('/health')
 def health():
-    return jsonify({
-        'status': 'healthy',
-        'device': str(DEVICE),
-        'models_loaded': len(model_cache.models),
-        'gpu_available': torch.cuda.is_available()
-    })
+    return jsonify({'status': 'healthy', 'models_loaded': len(cache.models)})
 
 @app.route('/api/category', methods=['POST'])
 def predict_category():
     try:
-        start_time = time.time()
-        data = request.get_json()
-        if not data or 'query' not in data:
-            return jsonify({'error': 'Missing "query" in request body'}), 400
-        
-        query = data['query']
-        model, tokenizer = model_cache.get_category_model()
-        
-        input_text = f"categorize: {query}"
-        input_ids = tokenizer.encode(input_text, return_tensors="pt").to(DEVICE)
-        
+        data = request.json
+        query = data.get('query', '')
+        model, tokenizer = cache.get_category_model()
+        input_ids = tokenizer.encode(f"categorize: {query}", return_tensors="pt").to(DEVICE)
         with torch.no_grad():
-            outputs = model.generate(input_ids, max_length=32, num_beams=4, early_stopping=True)
-        
+            outputs = model.generate(input_ids, max_length=32)
         category = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        inference_time = time.time() - start_time
-        
-        return jsonify({
-            'query': query,
-            'category': category,
-            'inference_time': f"{inference_time:.3f}s"
-        })
+        return jsonify({'query': query, 'category': category})
     except Exception as e:
-        logger.error(f"Error in predict_category: {str(e)}")
         return jsonify({'error': str(e)}), 500
 
 @app.route('/api/chat', methods=['POST'])
 def answer_question():
     try:
-        start_time = time.time()
-        data = request.get_json()
-        if not data or 'question' not in data:
-            return jsonify({'error': 'Missing "question" in request body'}), 400
-        
-        question = data['question']
-        model, tokenizer = model_cache.get_chat_model()
-        
-        input_text = f"answer question: {question}"
-        input_ids = tokenizer.encode(input_text, return_tensors="pt").to(DEVICE)
-        
+        data = request.json
+        question = data.get('question', '')
+        model, tokenizer = cache.get_chat_model()
+        input_ids = tokenizer.encode(f"answer question: {question}", return_tensors="pt").to(DEVICE)
         with torch.no_grad():
-            outputs = model.generate(input_ids, max_length=256, num_beams=4, early_stopping=True)
-        
+            outputs = model.generate(input_ids, max_length=256)
         answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        inference_time = time.time() - start_time
-        
-        return jsonify({
-            'question': question,
-            'answer': answer,
-            'inference_time': f"{inference_time:.3f}s"
-        })
+        return jsonify({'question': question, 'answer': answer})
     except Exception as e:
-        logger.error(f"Error in answer_question: {str(e)}")
         return jsonify({'error': str(e)}), 500
 
 @app.route('/api/service', methods=['POST'])
 def generate_description():
     try:
-        start_time = time.time()
-        data = request.get_json()
-        if not data or 'service_info' not in data:
-            return jsonify({'error': 'Missing "service_info" in request body'}), 400
-        
-        service_info = data['service_info']
-        model, tokenizer = model_cache.get_service_model()
-        
-        input_text = f"generate professional description: {service_info}"
-        input_ids = tokenizer.encode(input_text, return_tensors="pt").to(DEVICE)
-        
+        data = request.json
+        service_info = data.get('service_info', '')
+        model, tokenizer = cache.get_service_model()
+        input_ids = tokenizer.encode(f"generate professional description: {service_info}", return_tensors="pt").to(DEVICE)
         with torch.no_grad():
-            outputs = model.generate(input_ids, max_length=512, num_beams=4, early_stopping=True)
-        
+            outputs = model.generate(input_ids, max_length=512)
         description = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        inference_time = time.time() - start_time
-        
-        return jsonify({
-            'service_info': service_info,
-            'description': description,
-            'inference_time': f"{inference_time:.3f}s"
-        })
+        return jsonify({'service_info': service_info, 'description': description})
     except Exception as e:
-        logger.error(f"Error in generate_description: {str(e)}")
         return jsonify({'error': str(e)}), 500
 
 @app.route('/api/search', methods=['POST'])
 def semantic_search():
     try:
-        start_time = time.time()
-        data = request.get_json()
-        if not data or 'query' not in data or 'documents' not in data:
-            return jsonify({'error': 'Missing "query" or "documents" in request body'}), 400
-        
-        query = data['query']
-        documents = data['documents']
-        
-        if not isinstance(documents, list):
-            return jsonify({'error': '"documents" must be a list'}), 400
-        
-        model = model_cache.get_search_model()
-        
-        query_embedding = model.encode([query])
-        doc_embeddings = model.encode(documents)
-        
-        similarities = cosine_similarity(query_embedding, doc_embeddings)[0]
-        
-        results = [
-            {
-                'document': doc,
-                'similarity': float(score),
-                'rank': i + 1
-            }
-            for i, (doc, score) in enumerate(
-                sorted(zip(documents, similarities), key=lambda x: x[1], reverse=True)
-            )
-        ]
-        
-        inference_time = time.time() - start_time
-        
-        return jsonify({
-            'query': query,
-            'results': results,
-            'inference_time': f"{inference_time:.3f}s"
-        })
+        data = request.json
+        query = data.get('query', '')
+        documents = data.get('documents', [])
+        model = cache.get_search_model()
+        query_emb = model.encode([query])
+        doc_embs = model.encode(documents)
+        sims = cosine_similarity(query_emb, doc_embs)[0]
+        results = [{'document': d, 'similarity': float(s), 'rank': i+1} for i, (d, s) in enumerate(sorted(zip(documents, sims), key=lambda x: x[1], reverse=True))]
+        return jsonify({'query': query, 'results': results})
     except Exception as e:
-        logger.error(f"Error in semantic_search: {str(e)}")
         return jsonify({'error': str(e)}), 500
 
 @app.route('/api/moderate', methods=['POST'])
 def moderate_content():
     try:
-        start_time = time.time()
-        data = request.get_json()
-        if not data or 'text' not in data:
-            return jsonify({'error': 'Missing "text" in request body'}), 400
-        
-        text = data['text']
-        model, tokenizer = model_cache.get_moderation_model()
-        
+        data = request.json
+        text = data.get('text', '')
+        model, tokenizer = cache.get_moderation_model()
         inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512)
         inputs = {k: v.to(DEVICE) for k, v in inputs.items()}
-        
         with torch.no_grad():
             outputs = model(**inputs)
-            probabilities = F.softmax(outputs.logits, dim=-1)
-            toxic_prob = probabilities[0][1].item()
-        
-        is_toxic = toxic_prob > 0.5
-        inference_time = time.time() - start_time
-        
-        return jsonify({
-            'text': text,
-            'is_toxic': is_toxic,
-            'toxicity_score': round(toxic_prob, 4),
-            'status': 'toxic' if is_toxic else 'safe',
-            'inference_time': f"{inference_time:.3f}s"
-        })
-    except Exception as e:
-        logger.error(f"Error in moderate_content: {str(e)}")
-        return jsonify({'error': str(e)}), 500
-
-@app.route('/api/batch', methods=['POST'])
-def batch_process():
-    try:
-        data = request.get_json()
-        if not data or 'requests' not in data:
-            return jsonify({'error': 'Missing "requests" in request body'}), 400
-        
-        results = []
-        for req in data['requests']:
-            req_type = req.get('type')
-            
-            if req_type == 'category':
-                model, tokenizer = model_cache.get_category_model()
-                input_text = f"categorize: {req['query']}"
-                input_ids = tokenizer.encode(input_text, return_tensors="pt").to(DEVICE)
-                with torch.no_grad():
-                    outputs = model.generate(input_ids, max_length=32)
-                result = tokenizer.decode(outputs[0], skip_special_tokens=True)
-                results.append({'type': 'category', 'result': result})
-                
-            elif req_type == 'chat':
-                model, tokenizer = model_cache.get_chat_model()
-                input_text = f"answer question: {req['question']}"
-                input_ids = tokenizer.encode(input_text, return_tensors="pt").to(DEVICE)
-                with torch.no_grad():
-                    outputs = model.generate(input_ids, max_length=256)
-                result = tokenizer.decode(outputs[0], skip_special_tokens=True)
-                results.append({'type': 'chat', 'result': result})
-        
-        return jsonify({'results': results})
+            probs = F.softmax(outputs.logits, dim=-1)
+            toxic_prob = probs[0][1].item()
+        return jsonify({'text': text, 'is_toxic': toxic_prob > 0.5, 'toxicity_score': round(toxic_prob, 4), 'status': 'toxic' if toxic_prob > 0.5 else 'safe'})
     except Exception as e:
-        logger.error(f"Error in batch_process: {str(e)}")
         return jsonify({'error': str(e)}), 500
 
-@app.errorhandler(404)
-def not_found(error):
-    return jsonify({'error': 'Endpoint not found'}), 404
-
-@app.errorhandler(500)
-def internal_error(error):
-    return jsonify({'error': 'Internal server error'}), 500
-
 if __name__ == '__main__':
-    port = int(os.getenv('PORT', 7860))
-    logger.info(f"Starting server on port {port}...")
-    app.run(host='0.0.0.0', port=port, debug=False)
+    app.run(host='0.0.0.0', port=int(os.getenv('PORT', 7860)))