cai / app.py
msamogh's picture
Add lookup tables to NLU model
721b99b
raw
history blame
2.13 kB
from pprint import pprint, pformat
import os
import gradio as gr
import click
import torch
from rasa.nlu.model import Interpreter
from transformers import AutoTokenizer, AutoModelForSequenceClassification
# Rasa intent + entity extractor
RASA_MODEL_PATH = "woz_nlu_agent/models/nlu-lookup-1"
interpreter = Interpreter.load(RASA_MODEL_PATH)
# OOS classifier
OOS_MODEL= "msamogh/autonlp-cai-out-of-scope-649919116"
tokenizer = AutoTokenizer.from_pretrained(OOS_MODEL)
model = AutoModelForSequenceClassification.from_pretrained(OOS_MODEL)
MODEL_TYPES = {
"Out-of-scope classifier": "oos",
"Intent classifier": "intent_transformer",
"Intent and Entity extractor": "rasa_intent_entity"
}
def predict(model_type, input):
if MODEL_TYPES[model_type] == "rasa_intent_entity":
return rasa_predict(input)
elif MODEL_TYPES[model_type] == "oos":
return oos_predict(input)
elif MODEL_TYPES[model_type] == "intent_transformer":
return "WIP: intent_transformer"
def oos_predict(input):
inputs = tokenizer(input, return_tensors="pt")
outputs = model(**inputs).logits
outputs = torch.softmax(torch.tensor(outputs), dim=-1)[0]
return str({"In scope": outputs[1], "Out of scope": outputs[0]})
def rasa_predict(input):
def rasa_output(text):
message = str(text).strip()
result = interpreter.parse(message)
return result
response = rasa_output(input)
del response["response_selector"]
response["intent_ranking"] = response["intent_ranking"][:3]
if "id" in response["intent"]:
del response["intent"]["id"]
for i in response["intent_ranking"]:
if "id" in i:
del i["id"]
for e in response["entities"]:
if "extractor" in e:
del e["extractor"]
if "start" in e and "end" in e:
del e["start"]
del e["end"]
return pformat(response, indent=4)
def main():
iface = gr.Interface(fn=predict, inputs=[gr.inputs.Dropdown(list(MODEL_TYPES.keys())), "text"], outputs="text")
iface.launch()
if __name__ == "__main__":
main()