llmAPI.py

from flask import Flask, jsonify, request
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, StoppingCriteria, StoppingCriteriaList

print("Starting model loader...")

tokenizer = AutoTokenizer.from_pretrained("stabilityai/stablelm-tuned-alpha-3b")
model = AutoModelForCausalLM.from_pretrained("stabilityai/stablelm-tuned-alpha-3b")
model.half().cuda()

class StopOnTokens(StoppingCriteria):
    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
        stop_ids = [50278, 50279, 50277, 1, 0]
        for stop_id in stop_ids:
            if input_ids[0][-1] == stop_id:
                return True
        return False
    
def remove_text_before_string(text, input_string):
    index = text.find(input_string)
    if index == -1:
        return text
    else:
        return text[index + len(input_string):]


app = Flask(__name__)

@app.route("/llm", methods=["POST"])
def LLM():
    req_data = request.json
    inputPrompt = req_data['prompt']

    system_prompt = """<|SYSTEM|># StableLM Tuned (Alpha version)
    - StableLM is a helpful and harmless open-source AI language model developed by StabilityAI.
    - StableLM is excited to be able to help the user, but will refuse to do anything that could be considered harmful to the user.
    - StableLM is more than just an information source, StableLM is also able to write poetry, short stories, and make jokes.
    - StableLM will refuse to participate in anything that could harm a human.
    """


    prompt = f"{system_prompt}<|USER|>{inputPrompt}<|ASSISTANT|>"

    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
    tokens = model.generate(
    **inputs,
    max_new_tokens=20000,
    temperature=0.7,
    do_sample=True,
    stopping_criteria=StoppingCriteriaList([StopOnTokens()])
    )

    result = tokenizer.decode(tokens[0], skip_special_tokens=True)

    return jsonify({'result': remove_text_before_string(result, inputPrompt)})


if __name__ == "__main__":
    print("api available on port: 5555")
    print("/llm - {prompt: ""}")
    app.run(port=5555)