Upload 3 files

app/main.py

@@ -0,0 +1,26 @@
+from fastapi import FastAPI
+from pydantic import BaseModel
+from model.model import LLM
+import torch
+
+app = FastAPI()
+
+class InputText(BaseModel):
+    text: str
+
+# "bigscience/bloomz-1b1"
+model_tag = "facebook/opt-125m"
+model = LLM(model_name = model_tag,
+            device = "cuda" if torch.cuda.is_available() else "cpu")
+
+@app.get("/")
+async def root():
+    return {"message": "Technical challenge OK"}
+
+@app.post("/language-detection")
+def language_detection(text):
+    return {"language": model.language_detection(text)}
+
+@app.post("/entity-recognition")
+def ner(text):
+    return model.entity_recognition(text)

app/model/model.py

@@ -0,0 +1,77 @@
+import re
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+
+class LLM:
+    def __init__(self, model_name, device="cpu"):
+        # Model and tokenizer initialization
+        self.model, self.tokenizer = self.load_model_and_tokenizer(model_name, device)
+        # BCP-47 codes for the 3 available languages + unknown language
+        self.lang_codes = {
+            "english": "en",
+            "español": "es",
+            "française": "fr",
+            "unknown": "unk"}
+        
+    def load_model_and_tokenizer(self, model_name, device):
+        # Configuration for quantization (only works on GPU)
+        bnb_config = BitsAndBytesConfig(
+            use_4bit=True,
+            bnb_4bit_compute_dtype="float16",
+            bnb_4bit_quant_type="nf4",
+            use_nested_quant=False,
+        )
+        # Load model and tokenizer
+        model = AutoModelForCausalLM.from_pretrained(
+            model_name,
+            quantization_config=bnb_config,
+            cache_dir="./model_dir"
+        ).to(device)
+        tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True, cache_dir="./model_dir")
+        if tokenizer.pad_token_id is None:
+            tokenizer.pad_token_id = tokenizer.eos_token_id
+        return model, tokenizer
+
+    def language_detection(self, input_text):
+        # Prompt with one shot for each language
+        prompt = f"""Identify the language of the following sentences. Options: 'english', 'español', 'française' .
+            * <Identity theft is not a joke, millions of families suffer every year>(english)
+            * <Paseo a mi perro por el parque>(español)         
+            * <J'ai vu trop de souris à Paris>(française)
+            * <{input_text}>"""
+        # Generation and extraction of the language tag
+        answer_ids = self.model.generate(**self.tokenizer([prompt], return_tensors="pt"), max_new_tokens=10)
+        answer = self.tokenizer.batch_decode(answer_ids, skip_special_tokens=False)[0].split(prompt)[1]
+        pattern = r'\b(?:' + '|'.join(map(re.escape, self.lang_codes.keys())) + r')\b'
+        lang = re.search(pattern, answer, flags=re.IGNORECASE)
+        # Returns tag identified or 'unk' if none is detected
+        return self.lang_codes[lang.group()] if lang else self.lang_codes["unknown"]
+    
+    def entity_recognition(self, input_text):
+        # Prompt design
+        prompt = f"""Identify NER tags of 'location', 'organization', 'person' in the text.
+        
+        * Text: I saw Carmelo Anthony before the Knicks game in New York. Carmelo Anthony is retired now
+        * Tags: <Carmelo Anthony>(person), <Knicks>(organization), <New York>(location), <Carmelo Anthony>(person)
+        
+        * Text: I will work from Spain for LanguageWire because Spain is warmer than Denmark
+        * Tags: <Spain>(location), <LanguageWire>(organization), <Spain>(location), <Denmark>(location)
+        
+        * Text: Tesla founder Elon Musk is so rich that bought Twitter just for fun
+        * Tags: <Tesla>(organization), <Elon Musk>(person), <Twitter>(organization)
+        
+        * Text: {input_text}
+        * Tags: """
+        print(prompt)
+        # Generation and extraction of the identified entities
+        answer_ids = self.model.generate(**self.tokenizer([prompt], return_tensors="pt"), max_new_tokens=100)
+        answer = self.tokenizer.batch_decode(answer_ids, skip_special_tokens=True)[0].split(prompt)[1]
+        entities = re.findall(r'<(.*?)>', answer)
+        # Count of the tags detected (ignoring the type of entity)
+        entities_count = {}
+        for entity in entities:
+            if entity in entities_count:
+                entities_count[entity] += 1
+            else:
+                entities_count[entity] = 1
+        # Returns a dictionary
+        return entities_count