from flask import Flask, request, Response, jsonify from huggingface_hub import InferenceClient client = InferenceClient("mistralai/Mixtral-8x7B-Instruct-v0.1") def format_prompt(message, history): prompt = "" for user_prompt, bot_response in history: prompt += f"[INST] {user_prompt} [/INST]" prompt += f" {bot_response} " prompt += f"[INST] {message} [/INST]" return prompt def generate( prompt, history=[], temperature=0.2, max_new_tokens=2000, top_p=0.95, repetition_penalty=1.0, ): temperature = float(temperature) if temperature < 1e-2: temperature = 1e-2 top_p = float(top_p) generate_kwargs = dict( temperature=temperature, max_new_tokens=max_new_tokens, top_p=top_p, repetition_penalty=repetition_penalty, do_sample=True, seed=42, ) #formatted_prompt = format_prompt(prompt, history) #stream = client.text_generation(prompt, **generate_kwargs, stream=True, details=False, return_full_text=False) response = client.text_generation(prompt, **generate_kwargs, stream=False, details=False, return_full_text=False) print(response) return response.encode('utf-8') #output = "" #for response in stream: # yield response.token.text.encode('utf-8') app = Flask(__name__) @app.route('/health', methods=['GET']) def health(): return jsonify({"status": "ok"}) @app.route('/completion', methods=['POST']) def search_route(): data = request.get_json() prompt = data.get('prompt', '') #truncated_prompt = prompt[:32768] return Response(generate(prompt), content_type='text/plain; charset=utf-8', status=200, direct_passthrough=True) if __name__ == '__main__': app.run(debug=False, host='0.0.0.0', port=7860)