fix squeeze issue + load only once pipeline

app.py

@@ -4,9 +4,10 @@ from fastapi import FastAPI, UploadFile
 from typing import Union 
 import json 
 import csv
-from modeles import bert, squeezebert, deberta
+from modeles import bert, squeezebert, deberta, loadSqueeze
 from uploadFile import file_to_text
 from typing import List
+from transformers import pipeline
 
 
 
@@ -21,10 +22,10 @@ app.add_middleware(
 )
 
 
-@app.on_event("startup")
-async def startup_event():
-    print("start")
 
+pipBert = pipeline('question-answering', model="ALOQAS/bert-large-uncased-finetuned-squad-v2", tokenizer="ALOQAS/bert-large-uncased-finetuned-squad-v2")
+pipDeberta = pipeline('question-answering', model="ALOQAS/deberta-large-finetuned-squad-v2", tokenizer="ALOQAS/deberta-large-finetuned-squad-v2")
+tokenizer, model = loadSqueeze()
 
 @app.get("/")
 async def root():
@@ -60,7 +61,7 @@ async def create_upload_file(texte: str, model: str):
 @app.post("/squeezebert/")
 async def qasqueezebert(context: str, question: str):
     try:
-        squeezebert_answer = squeezebert(context, question)
+        squeezebert_answer = squeezebert(context, question, model, tokenizer)
         if squeezebert_answer:
             return squeezebert_answer
         else:
@@ -75,7 +76,7 @@ async def qasqueezebert(context: str, question: str):
 @app.post("/bert/")
 async def qabert(context: str, question: str):
     try:
-        bert_answer = bert(context, question)
+        bert_answer = bert(context, question, pipBert)
         if bert_answer:
             return bert_answer
         else:
@@ -90,7 +91,7 @@ async def qabert(context: str, question: str):
 @app.post("/deberta-v2/")
 async def qadeberta(context: str, question: str):
     try:
-        deberta_answer = deberta(context, question)
+        deberta_answer = deberta(context, question, pipDeberta)
         if deberta_answer:
             return deberta_answer
         else:

modeles.py

@@ -1,19 +1,37 @@
-from transformers import pipeline
+from transformers import AutoTokenizer, AutoModelForQuestionAnswering
+import torch
 
-def squeezebert(context, question):
+def loadSqueeze():
+    tokenizer = AutoTokenizer.from_pretrained("ALOQAS/squeezebert-uncased-finetuned-squad-v2")
+    model = AutoModelForQuestionAnswering.from_pretrained("ALOQAS/squeezebert-uncased-finetuned-squad-v2")
+    return tokenizer, model
+
+def squeezebert(context, question, model, tokenizer):
     # Define the specific model and tokenizer for SqueezeBERT
-    model_name = "ALOQAS/squeezebert-uncased-finetuned-squad-v2"
-    pip = pipeline('question-answering', model=model_name, tokenizer=model_name)
-    return pip(context=context, question=question)
+        # Tokenize the input question-context pair
+    inputs = tokenizer.encode_plus(question, context, max_length=512, truncation=True, padding=True, return_tensors='pt') 
+
+    # Send inputs to the same device as your model
+    inputs = {k: v.to(model.device) for k, v in inputs.items()}
+
+    with torch.no_grad():
+        # Forward pass, get model outputs
+        outputs = model(**inputs)
+
+    # Extract the start and end positions of the answer in the tokens
+    answer_start_scores, answer_end_scores = outputs.start_logits, outputs.end_logits
+    answer_start_index = torch.argmax(answer_start_scores)  # Most likely start of answer
+    answer_end_index = torch.argmax(answer_end_scores) + 1  # Most likely end of answer; +1 for inclusive slicing
+
+    # Convert token indices to the actual answer text
+    answer_tokens = inputs['input_ids'][0, answer_start_index:answer_end_index]
+    answer = tokenizer.decode(answer_tokens, skip_special_tokens=True)
+    return {"answer": answer, "start": answer_start_index.item(), "end": answer_end_index.item()}
+
+
 
-def bert(context, question):
-    # Define the specific model and tokenizer for BERT
-    model_name = "ALOQAS/bert-large-uncased-finetuned-squad-v2"
-    pip = pipeline('question-answering', model=model_name, tokenizer=model_name)
+def bert(context, question, pip):
     return pip(context=context, question=question)
 
-def deberta(context, question):
-    # Define the specific model and tokenizer for DeBERTa
-    model_name = "ALOQAS/deberta-large-finetuned-squad-v2"
-    pip = pipeline('question-answering', model=model_name, tokenizer=model_name)
+def deberta(context, question, pip):
     return pip(context=context, question=question)