Upload main.py

main.py

@@ -1,11 +1,11 @@
 import torch
 from fastapi import FastAPI, Depends, status
 from fastapi.responses import PlainTextResponse
-from transformers import AutoTokenizer, AutoModel, DPRQuestionEncoder
+from pydantic import BaseModel
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
-from datasets import load_from_disk
 import time
-from typing import Dict
+from typing import Dict, List, Optional
 
 import jwt
 from decouple import config
@@ -17,22 +17,12 @@ JWT_ALGORITHM = config("algorithm")
 
 app = FastAPI()
 app.ready = False
-columns = ['kilt_id', 'wikipedia_id', 'wikipedia_title', 'text', 'anchors', 'categories',
-           'wikidata_info', 'history']
 
-min_snippet_length = 20
-topk = 21
 device = ("cuda" if torch.cuda.is_available() else "cpu")
-model = DPRQuestionEncoder.from_pretrained("vblagoje/dpr-question_encoder-single-lfqa-wiki").to(device)
-tokenizer = AutoTokenizer.from_pretrained("vblagoje/dpr-question_encoder-single-lfqa-wiki")
+tokenizer = AutoTokenizer.from_pretrained('vblagoje/bart_lfqa')
+model = AutoModelForSeq2SeqLM.from_pretrained('vblagoje/bart_lfqa').to(device)
 _ = model.eval()
 
-index_file_name = "./data/kilt_wikipedia.faiss"
-
-kilt_wikipedia_paragraphs = load_from_disk("./data/kilt_wiki_prepared")
-# use paragraphs that are not simple fragments or very short sentences
-kilt_wikipedia_paragraphs = kilt_wikipedia_paragraphs.filter(lambda x: x["end_character"] > 200)
-
 
 class JWTBearer(HTTPBearer):
     def __init__(self, auto_error: bool = True):
@@ -85,27 +75,26 @@ def decodeJWT(token: str) -> dict:
         return {}
 
 
-def embed_questions_for_retrieval(questions):
-    query = tokenizer(questions, max_length=128, padding=True, truncation=True, return_tensors="pt")
-    with torch.no_grad():
-        q_reps = model(query["input_ids"].to(device), query["attention_mask"].to(device)).pooler_output
-    return q_reps.cpu().numpy()
+class LFQAParameters(BaseModel):
+    min_length: int = 50
+    max_length: int = 250
+    do_sample: bool = False
+    early_stopping: bool = True
+    num_beams: int = 8
+    temperature: float = 1.0
+    top_k: float = None
+    top_p: float = None
+    no_repeat_ngram_size: int = 3
+    num_return_sequences: int = 1
 
-def query_index(question):
-    question_embedding = embed_questions_for_retrieval([question])
-    scores, wiki_passages = kilt_wikipedia_paragraphs.get_nearest_examples("embeddings", question_embedding, k=topk)
-    columns = ['wikipedia_id', 'title', 'text', 'section', 'start_paragraph_id', 'end_paragraph_id',
-               'start_character', 'end_character']
-    retrieved_examples = []
-    r = list(zip(wiki_passages[k] for k in columns))
-    for i in range(topk):
-        retrieved_examples.append({k: v for k, v in zip(columns, [r[j][0][i] for j in range(len(columns))])})
-    return retrieved_examples
+
+class InferencePayload(BaseModel):
+    model_input: str
+    parameters: Optional[LFQAParameters] = LFQAParameters()
 
 
 @app.on_event("startup")
 def startup():
-    kilt_wikipedia_paragraphs.load_faiss_index("embeddings", index_file_name, device=0)
     app.ready = True
 
 
@@ -116,7 +105,26 @@ def healthz():
     return PlainTextResponse("service unavailable", status_code=status.HTTP_503_SERVICE_UNAVAILABLE)
 
 
-@app.get("/find_context", dependencies=[Depends(JWTBearer())])
-def find_context(question: str = None):
-    return [res for res in query_index(question) if len(res["text"].split()) > min_snippet_length][:int(topk / 3)]
-
+@app.post("/generate/", dependencies=[Depends(JWTBearer())])
+def generate(context: InferencePayload):
+
+    model_input = tokenizer(context.model_input, truncation=True, padding=True, return_tensors="pt")
+    param = context.parameters
+    generated_answers_encoded = model.generate(input_ids=model_input["input_ids"].to(device),
+                                               attention_mask=model_input["attention_mask"].to(device),
+                                               min_length=param.min_length,
+                                               max_length=param.max_length,
+                                               do_sample=param.do_sample,
+                                               early_stopping=param.early_stopping,
+                                               num_beams=param.num_beams,
+                                               temperature=param.temperature,
+                                               top_k=param.top_k,
+                                               top_p=param.top_p,
+                                               no_repeat_ngram_size=param.no_repeat_ngram_size,
+                                               num_return_sequences=param.num_return_sequences)
+    answers = tokenizer.batch_decode(generated_answers_encoded, skip_special_tokens=True,
+                                     clean_up_tokenization_spaces=True)
+    results = []
+    for answer in answers:
+        results.append({"generated_text": answer})
+    return results