Update api.py

api.py

@@ -1,38 +1,59 @@
-from flask import Flask, request, jsonify
-from transformers import AutoTokenizer, AutoModelForCausalLM
-import torch
+from flask import Flask, request, jsonify, Response
+import subprocess
+import json
+import logging
 
-app = Flask(__name__)
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
 
-# Load model and tokenizer (outside the route)
-model_name = "gpt2"  # Replace with your desired model name (e.g., gpt-neo)
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name).to("cuda" if torch.cuda.is_available() else "cpu")
+app = Flask(__name__)
 
-@app.route("/completions", methods=["POST"])
-def generate_text():
+@app.route('/completions', methods=['POST'])
+def get_completion():
     try:
         data = request.get_json()
-        if not data or "prompt" not in data:
+        if not data or 'prompt' not in data:
             return jsonify({"error": "Missing 'prompt' in request body"}), 400
 
-        prompt = data["prompt"]
-        input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(model.device)
-
-        # Generate text in a loop, simulating streaming
-        max_length = 150
-        generated_tokens = []
-        for _ in range(max_length):
-            output = model.generate(input_ids, max_length=1, do_sample=True)
-            generated_tokens.append(output[0][0].item())
-            input_ids = torch.cat([input_ids, output], dim=1)
-
-        text_output = tokenizer.decode(generated_tokens, skip_special_tokens=True)
-        return jsonify({"text": text_output})
+        prompt = data['prompt']
+
+        def generate():
+            try:
+                process = subprocess.Popen(
+                    ["ollama", "run", "llama2", "-p", prompt],  # Replace "llama2" with your model
+                    stdout=subprocess.PIPE,
+                    stderr=subprocess.PIPE,
+                    text=True,
+                    encoding="utf-8"
+                )
+
+                for line in process.stdout:
+                    try:
+                        json_line = json.loads(line)
+                        if "response" in json_line:
+                            yield f"data: {json.dumps({'text': json_line['response']})}\n\n"
+                        elif "done" in json_line: #check for done message
+                            break
+                    except json.JSONDecodeError:
+                        logger.warning(f"Invalid JSON line from Ollama: {line.strip()}")
+
+                stderr = process.stderr.read()
+                if stderr:
+                    logger.error(f"Ollama stderr: {stderr}")
+
+                process.wait()
+
+            except FileNotFoundError:
+                yield f"data: {json.dumps({'text': 'Error: Ollama not found. Is it installed and in your PATH?'})}\n\n"
+            except Exception as e:
+                logger.exception("Error in Ollama subprocess:")
+                yield f"data: {json.dumps({'text': f'Error in Ollama: {e}'})}\n\n"
+
+        return Response(generate(), mimetype='text/event-stream')
 
     except Exception as e:
-        print(f"Error during generation: {e}")
-        return jsonify({"error": "An error occurred."}), 500
+        logger.exception("Error in /completions route:")
+        return jsonify({"error": "An error occurred during processing."}), 500
 
 if __name__ == "__main__":
-    app.run(debug=False, host="0.0.0.0", port=5000)
+    app.run(debug=False, host='0.0.0.0', port=5000)