Spaces:
Sleeping
Sleeping
Upload 3 files
Browse files- app/main.py +26 -0
- app/model/__pycache__/model.cpython-310.pyc +0 -0
- app/model/model.py +77 -0
app/main.py
ADDED
@@ -0,0 +1,26 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from fastapi import FastAPI
|
2 |
+
from pydantic import BaseModel
|
3 |
+
from model.model import LLM
|
4 |
+
import torch
|
5 |
+
|
6 |
+
app = FastAPI()
|
7 |
+
|
8 |
+
class InputText(BaseModel):
|
9 |
+
text: str
|
10 |
+
|
11 |
+
# "bigscience/bloomz-1b1"
|
12 |
+
model_tag = "facebook/opt-125m"
|
13 |
+
model = LLM(model_name = model_tag,
|
14 |
+
device = "cuda" if torch.cuda.is_available() else "cpu")
|
15 |
+
|
16 |
+
@app.get("/")
|
17 |
+
async def root():
|
18 |
+
return {"message": "Technical challenge OK"}
|
19 |
+
|
20 |
+
@app.post("/language-detection")
|
21 |
+
def language_detection(text):
|
22 |
+
return {"language": model.language_detection(text)}
|
23 |
+
|
24 |
+
@app.post("/entity-recognition")
|
25 |
+
def ner(text):
|
26 |
+
return model.entity_recognition(text)
|
app/model/__pycache__/model.cpython-310.pyc
ADDED
Binary file (3.32 kB). View file
|
|
app/model/model.py
ADDED
@@ -0,0 +1,77 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import re
|
2 |
+
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
|
3 |
+
|
4 |
+
class LLM:
|
5 |
+
def __init__(self, model_name, device="cpu"):
|
6 |
+
# Model and tokenizer initialization
|
7 |
+
self.model, self.tokenizer = self.load_model_and_tokenizer(model_name, device)
|
8 |
+
# BCP-47 codes for the 3 available languages + unknown language
|
9 |
+
self.lang_codes = {
|
10 |
+
"english": "en",
|
11 |
+
"español": "es",
|
12 |
+
"française": "fr",
|
13 |
+
"unknown": "unk"}
|
14 |
+
|
15 |
+
def load_model_and_tokenizer(self, model_name, device):
|
16 |
+
# Configuration for quantization (only works on GPU)
|
17 |
+
bnb_config = BitsAndBytesConfig(
|
18 |
+
use_4bit=True,
|
19 |
+
bnb_4bit_compute_dtype="float16",
|
20 |
+
bnb_4bit_quant_type="nf4",
|
21 |
+
use_nested_quant=False,
|
22 |
+
)
|
23 |
+
# Load model and tokenizer
|
24 |
+
model = AutoModelForCausalLM.from_pretrained(
|
25 |
+
model_name,
|
26 |
+
quantization_config=bnb_config,
|
27 |
+
cache_dir="./model_dir"
|
28 |
+
).to(device)
|
29 |
+
tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True, cache_dir="./model_dir")
|
30 |
+
if tokenizer.pad_token_id is None:
|
31 |
+
tokenizer.pad_token_id = tokenizer.eos_token_id
|
32 |
+
return model, tokenizer
|
33 |
+
|
34 |
+
def language_detection(self, input_text):
|
35 |
+
# Prompt with one shot for each language
|
36 |
+
prompt = f"""Identify the language of the following sentences. Options: 'english', 'español', 'française' .
|
37 |
+
* <Identity theft is not a joke, millions of families suffer every year>(english)
|
38 |
+
* <Paseo a mi perro por el parque>(español)
|
39 |
+
* <J'ai vu trop de souris à Paris>(française)
|
40 |
+
* <{input_text}>"""
|
41 |
+
# Generation and extraction of the language tag
|
42 |
+
answer_ids = self.model.generate(**self.tokenizer([prompt], return_tensors="pt"), max_new_tokens=10)
|
43 |
+
answer = self.tokenizer.batch_decode(answer_ids, skip_special_tokens=False)[0].split(prompt)[1]
|
44 |
+
pattern = r'\b(?:' + '|'.join(map(re.escape, self.lang_codes.keys())) + r')\b'
|
45 |
+
lang = re.search(pattern, answer, flags=re.IGNORECASE)
|
46 |
+
# Returns tag identified or 'unk' if none is detected
|
47 |
+
return self.lang_codes[lang.group()] if lang else self.lang_codes["unknown"]
|
48 |
+
|
49 |
+
def entity_recognition(self, input_text):
|
50 |
+
# Prompt design
|
51 |
+
prompt = f"""Identify NER tags of 'location', 'organization', 'person' in the text.
|
52 |
+
|
53 |
+
* Text: I saw Carmelo Anthony before the Knicks game in New York. Carmelo Anthony is retired now
|
54 |
+
* Tags: <Carmelo Anthony>(person), <Knicks>(organization), <New York>(location), <Carmelo Anthony>(person)
|
55 |
+
|
56 |
+
* Text: I will work from Spain for LanguageWire because Spain is warmer than Denmark
|
57 |
+
* Tags: <Spain>(location), <LanguageWire>(organization), <Spain>(location), <Denmark>(location)
|
58 |
+
|
59 |
+
* Text: Tesla founder Elon Musk is so rich that bought Twitter just for fun
|
60 |
+
* Tags: <Tesla>(organization), <Elon Musk>(person), <Twitter>(organization)
|
61 |
+
|
62 |
+
* Text: {input_text}
|
63 |
+
* Tags: """
|
64 |
+
print(prompt)
|
65 |
+
# Generation and extraction of the identified entities
|
66 |
+
answer_ids = self.model.generate(**self.tokenizer([prompt], return_tensors="pt"), max_new_tokens=100)
|
67 |
+
answer = self.tokenizer.batch_decode(answer_ids, skip_special_tokens=True)[0].split(prompt)[1]
|
68 |
+
entities = re.findall(r'<(.*?)>', answer)
|
69 |
+
# Count of the tags detected (ignoring the type of entity)
|
70 |
+
entities_count = {}
|
71 |
+
for entity in entities:
|
72 |
+
if entity in entities_count:
|
73 |
+
entities_count[entity] += 1
|
74 |
+
else:
|
75 |
+
entities_count[entity] = 1
|
76 |
+
# Returns a dictionary
|
77 |
+
return entities_count
|