sjaipuriyar/phi3_model_testing
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Fixed some routing issues, and allowed support for multiple models

Browse source
This commit is contained in:
fedora Cloud User 2024-08-22 23:16:27 +00:00
parent 6806160f46
commit de1f707f1e
1 changed files with 17 additions and 5 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

22
testing.py
View file

@ -5,12 +5,24 @@ from flask import Flask, request
app = Flask(__name__) app = Flask(__name__)
@app.route("/query",methods=["POST"])
def generate_response(): @app.route("/query/",methods=["POST"])
@app.route("/query/<string:req_model>",methods=["POST"])
def generate_response(req_model="phi"):
if req_model not in ['phi','granite']: return "Only models phi and granite are supported.", 400
try: try:
query = request.form['query'] query = request.form['query']
if req_model == 'granite':
model = model_granite
tokenizer = tokenizer_granite
else:
model = model_phi
tokenizer = tokenizer_phi
messages = [ messages = [
{"role": "system", "content": "You are an evil and unhelpful AI assistant."}, {"role": "system", "content": "You are a helpful AI assistant."},
{"role": "user","content": query} {"role": "user","content": query}
] ]
pipe = pipeline( pipe = pipeline(
@ -34,13 +46,13 @@ def generate_response():
if __name__ == '__main__': if __name__ == '__main__':
torch.random.manual_seed(0) torch.random.manual_seed(0)
model = AutoModelForCausalLM.from_pretrained( model_phi = AutoModelForCausalLM.from_pretrained(
"/home/fedora/microsoft/Phi-3-mini-4k-instruct", "/home/fedora/microsoft/Phi-3-mini-4k-instruct",
device_map="cuda", device_map="cuda",
torch_dtype="auto", torch_dtype="auto",
trust_remote_code=True, trust_remote_code=True,
) )
model = AutoModelForCausalLM.from_pretrained( model_granite = AutoModelForCausalLM.from_pretrained(
"/home/fedora/granite-3b-code-instruct", "/home/fedora/granite-3b-code-instruct",
device_map="cuda", device_map="cuda",
torch_dtype="auto", torch_dtype="auto",