Add app.py

app.py

@@ -0,0 +1,325 @@
+#!/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 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)
+logger = logging.getLogger(__name__)
+
+app = Flask(__name__)
+CORS(app)
+
+HF_USERNAME = os.getenv("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")
+        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")
+        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")
+        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")
+        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")
+        return self.models['moderation'], self.tokenizers['moderation']
+
+model_cache = ModelCache()
+
+@app.route('/', methods=['GET'])
+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'
+    })
+
+@app.route('/health', methods=['GET'])
+def health():
+    return jsonify({
+        'status': 'healthy',
+        'device': str(DEVICE),
+        'models_loaded': len(model_cache.models),
+        'gpu_available': torch.cuda.is_available()
+    })
+
+@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)
+        
+        with torch.no_grad():
+            outputs = model.generate(input_ids, max_length=32, num_beams=4, early_stopping=True)
+        
+        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"
+        })
+    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)
+        
+        with torch.no_grad():
+            outputs = model.generate(input_ids, max_length=256, num_beams=4, early_stopping=True)
+        
+        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"
+        })
+    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)
+        
+        with torch.no_grad():
+            outputs = model.generate(input_ids, max_length=512, num_beams=4, early_stopping=True)
+        
+        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"
+        })
+    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"
+        })
+    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()
+        
+        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})
+    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)