Spaces:
Sleeping
Sleeping
File size: 2,275 Bytes
f4f976d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 |
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from transformers import AutoModelForCausalLM, AutoTokenizer
import logging
app = FastAPI()
# Add logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
formatter = logging.Formatter("%(asctime)s - %(levelname)s - %(message)s")
handler = logging.StreamHandler()
handler.setFormatter(formatter)
logger.addHandler(handler)
# Add CORS
origins = ["*"]
app.add_middleware(
CORSMiddleware,
allow_origins=origins,
allow_credentials=True,
allow_methods=["GET", "POST", "PUT", "DELETE"],
allow_headers=["*"],
)
intent_model = AutoModelForCausalLM.from_pretrained("llmware/slim-intent")
intent_tokenizer = AutoTokenizer.from_pretrained("llmware/slim-intent")
sentiment_model = AutoModelForCausalLM.from_pretrained("llmware/slim-sentiment")
sentiment_tokenizer = AutoTokenizer.from_pretrained("llmware/slim-sentiment")
def getResponse(model, tokenizer, text, params):
function = "classify"
prompt = "<human>: " + text + "\n" + f"<{function}> {params} </{function}>\n<bot>:"
inputs = tokenizer(prompt, return_tensors="pt")
start_of_input = len(inputs.input_ids[0])
outputs = model.generate(
inputs.input_ids.to('cpu'),
eos_token_id=tokenizer.eos_token_id,
pad_token_id=tokenizer.eos_token_id,
do_sample=True,
temperature=0.3,
max_new_tokens=100
)
output = tokenizer.decode(outputs[0][start_of_input:], skip_special_tokens=True)
return output
@app.get("/")
def read_root():
return {
"message": "API running successfully",
"endpoints": [
"/api/sentiment/",
"/api/intent/"
]
}
@app.post("/api/intent/")
def intentResponse(text: str):
params = "intent"
try:
responses = getResponse(intent_model, intent_tokenizer, text, params)
return responses
except Exception as e:
logger.exception(e)
return {"API Error": str(e)}
@app.post("/api/sentiment/")
def sentimentResponse(text: str):
params = "sentiment"
try:
responses = getResponse(sentiment_model, sentiment_tokenizer, text, params)
return responses
except Exception as e:
logger.exception(e)
return {"API Error": str(e)} |