Update app.py

app.py

@@ -26,19 +26,20 @@ __license__ = "Apache 2.0"
 __version__ = "0.0.1"
 
 
-from flask import Flask, request, jsonify, render_template
+from flask import Flask, request, jsonify
 from flask_cors import CORS, cross_origin
 from flask_restful import Resource, Api, reqparse
 import control.recommendation_handler as recommendation_handler
-from helpers import get_credentials, authenticate_api, save_model
+from helpers import get_credentials, authenticate_api, save_model, inference
 import config as cfg
+import requests
 import logging
 import uuid
 import json
 import os
-import requests
+import pickle
 
-app = Flask(__name__, static_folder='static')
+app = Flask(__name__)
 
 # configure logging
 logging.basicConfig(
@@ -66,42 +67,60 @@ def index():
 @app.route("/recommend", methods=['GET'])
 @cross_origin()
 def recommend():
-    user_ip = request.remote_addr
-    hf_token, hf_url = get_credentials.get_credentials()
-    api_url, headers = authenticate_api.authenticate_api(hf_token, hf_url)
+    model_id, _ =save_model.save_model()
     prompt_json = recommendation_handler.populate_json()
     args = request.args
+    print("args list = ", args)
     prompt = args.get("prompt")
-    recommendation_json = recommendation_handler.recommend_prompt(prompt, prompt_json,
-                                                                  api_url, headers)
+
+    umap_model_file = './models/umap/sentence-transformers/all-MiniLM-L6-v2/umap.pkl'
+    with open(umap_model_file, 'rb') as f:
+        umap_model = pickle.load(f)
+
+    # Embeddings from HF API
+    # hf_token, hf_url = get_credentials.get_hf_credentials()
+    # api_url, headers = authenticate_api.authenticate_api(hf_token, hf_url)
+    # api_url = f'https://router.huggingface.co/hf-inference/models/{model_id}/pipeline/feature-extraction'
+    # embedding_fn = recommendation_handler.get_embedding_func(inference='huggingface', model_id=model_id, api_url= api_url, headers = headers)
+
+    # Embeddings from local inference
+    embedding_fn = recommendation_handler.get_embedding_func(inference='local', model_id=model_id)
+
+    recommendation_json = recommendation_handler.recommend_prompt(prompt, prompt_json, embedding_fn, umap_model=umap_model)
+
+    user_ip = request.remote_addr
     logger.info(f'USER - {user_ip} - ID {id} - accessed recommend route')
     logger.info(f'RECOMMEND ROUTE - request: {prompt} response: {recommendation_json}')
+
     return recommendation_json
 
 @app.route("/get_thresholds", methods=['GET'])
 @cross_origin()
 def get_thresholds():
-    hf_token, hf_url = get_credentials.get_credentials()
+    hf_token, hf_url = get_credentials.get_hf_credentials()
     api_url, headers = authenticate_api.authenticate_api(hf_token, hf_url)
     prompt_json = recommendation_handler.populate_json()
-    model_id = 'sentence-transformers/all-minilm-l6-v2'
     args = request.args
-    #print("args list = ", args)
     prompt = args.get("prompt")
-    thresholds_json = recommendation_handler.get_thresholds(prompt, prompt_json, api_url,
-                                                            headers, model_id)
+    thresholds_json = recommendation_handler.get_thresholds(prompt, prompt_json, api_url, headers)
     return thresholds_json
 
 @app.route("/recommend_local", methods=['GET'])
 @cross_origin()
 def recommend_local():
-    model_id, model_path = save_model.save_model()
-    prompt_json = recommendation_handler.populate_json()
+    model_id, _ = save_model.save_model()
+    prompt_json, _ = recommendation_handler.populate_json()
     args = request.args
     print("args list = ", args)
     prompt = args.get("prompt")
-    local_recommendation_json = recommendation_handler.recommend_local(prompt, prompt_json,
-                                                                       model_id, model_path)
+    
+    umap_model_file = './models/umap/sentence-transformers/all-MiniLM-L6-v2/umap.pkl'
+    with open(umap_model_file, 'rb') as f:
+        umap_model = pickle.load(f)
+
+    embedding_fn = recommendation_handler.get_embedding_func(inference='local', model_id=model_id)
+
+    local_recommendation_json = recommendation_handler.recommend_prompt(prompt, prompt_json, embedding_fn, umap_model=umap_model)
     return local_recommendation_json
 
 @app.route("/log", methods=['POST'])
@@ -127,51 +146,27 @@ def log():
 @cross_origin()
 def demo_inference():
     args = request.args
-    # model_id = "meta-llama/Llama-3.2-11B-Vision-Instruct"
-    model_id = args.get('model_id', default="meta-llama/Llama-4-Scout-17B-16E-Instruct")
-    temperature = 0.5
-    max_new_tokens = 1000
 
-    hf_token, hf_url = get_credentials.get_credentials()
+    inference_provider = args.get('inference_provider', default='replicate')
+    model_id = args.get('model_id', default="ibm-granite/granite-3.3-8b-instruct")
+    temperature = args.get('temperature', default=0.5)
+    max_new_tokens = args.get('max_new_tokens', default=1000)
 
     prompt = args.get('prompt')
 
-    API_URL = "https://router.huggingface.co/together/v1/chat/completions"
-    headers = {
-        "Authorization": f"Bearer {hf_token}",
-    }
-
-    response = requests.post(
-        API_URL,
-        headers=headers, 
-        json={
-            "messages": [
-                {
-                    "role": "user",
-                    "content": [
-                        {
-                            "type": "text",
-                            "text": prompt
-                        },
-                    ]
-                }
-            ],
-            "model": model_id,
-            'temperature': temperature,
-            'max_new_tokens': max_new_tokens,
-        }
-    )
     try:
-        response = response.json()["choices"][0]["message"]
+        response = inference.INFERENCE_HANDLER[inference_provider](prompt, model_id, temperature, max_new_tokens)
         response.update({
+            'inference_provider': inference_provider,
             'model_id': model_id,
             'temperature': temperature,
             'max_new_tokens': max_new_tokens,
         })
+
         return response
     except:
-        return response.text, response.status_code
+        return "Model Inference failed.", 500
 
 if __name__=='__main__':
-    debug_mode = os.getenv('FLASK_DEBUG', 'True').lower() in ['true', '1', 't']
-    app.run(host='0.0.0.0', port='7860', debug=debug_mode)
+    debug_mode = os.getenv('FLASK_DEBUG', 'False').lower() in ['true', '1', 't']
+    app.run(host='0.0.0.0', port='8080', debug=debug_mode)