| import os |
| import json |
| import requests as _requests |
| from bs4 import BeautifulSoup |
| from flask import Flask, render_template, request, jsonify, stream_with_context, Response |
| from airllm import AutoModel |
|
|
| app = Flask(__name__) |
|
|
| model = None |
| chat_history = [] |
|
|
| MODEL_NAME = "Exocore/ExocoreV1" |
|
|
|
|
| def load_model(): |
| global model |
| try: |
| model = AutoModel.from_pretrained(MODEL_NAME) |
| except Exception as e: |
| print(f"Model load error: {e}") |
|
|
|
|
| @app.route("/") |
| def index(): |
| return render_template("index.html") |
|
|
|
|
| @app.route("/api/model-details", methods=["GET"]) |
| def model_details(): |
| if model is None: |
| return jsonify({"error": "No model loaded"}), 400 |
| return jsonify({ |
| "name": "ExocoreV1", |
| "type": "airllm", |
| }) |
|
|
|
|
| @app.route("/api/chat", methods=["POST"]) |
| def chat(): |
| global chat_history |
| if model is None: |
| return jsonify({"error": "Model not loaded"}), 400 |
|
|
| data = request.json |
| user_message = data.get("message", "").strip() |
| max_tokens = int(data.get("max_tokens", 128)) |
|
|
| if not user_message: |
| return jsonify({"error": "Empty message"}), 400 |
|
|
| def generate(): |
| full_reply = "" |
| try: |
| input_text = [user_message] |
| input_tokens = model.tokenizer( |
| input_text, |
| return_tensors="pt", |
| return_attention_mask=False, |
| truncation=True, |
| max_length=128, |
| padding=False |
| ) |
| generation_output = model.generate( |
| input_tokens["input_ids"].cuda(), |
| max_new_tokens=max_tokens, |
| use_cache=True, |
| return_dict_in_generate=True |
| ) |
| output = model.tokenizer.decode(generation_output.sequences[0]) |
| full_reply = output |
| for ch in full_reply: |
| yield f"data: {json.dumps({'token': ch})}\n\n" |
| chat_history.append({"role": "user", "content": user_message}) |
| chat_history.append({"role": "assistant", "content": full_reply}) |
| yield f"data: {json.dumps({'done': True})}\n\n" |
| except Exception as e: |
| yield f"data: {json.dumps({'error': str(e)})}\n\n" |
|
|
| return Response(stream_with_context(generate()), mimetype="text/event-stream") |
|
|
|
|
| @app.route("/api/clear", methods=["POST"]) |
| def clear_history(): |
| global chat_history |
| chat_history = [] |
| return jsonify({"success": True}) |
|
|
|
|
| @app.route("/api/status", methods=["GET"]) |
| def status(): |
| return jsonify({ |
| "model_loaded": model is not None, |
| "history_length": len(chat_history), |
| }) |
|
|
|
|
| if __name__ == "__main__": |
| load_model() |
| port = int(os.environ.get("PORT", 7860)) |
| app.run(host="0.0.0.0", port=port, debug=False) |
|
|
