main

Dockerfile

@@ -0,0 +1,22 @@
+# Use the official Python image as the base
+FROM python:3.9-slim
+
+# Set environment variables
+ENV PYTHONUNBUFFERED=1 \
+    PYTHONDONTWRITEBYTECODE=1
+
+# Create a working directory
+WORKDIR /app
+
+# Copy requirements and install dependencies
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Copy the application code
+COPY . .
+
+# Expose the port FastAPI will run on
+EXPOSE 8000
+
+# Command to run the application
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "8000"]

app/main.py

@@ -0,0 +1,42 @@
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+from sentence_transformers import SentenceTransformer, util
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+import torch
+
+app = FastAPI()
+
+# Load the retriever model
+retriever = SentenceTransformer('all-MiniLM-L6-v2')
+
+# Load the generator model
+tokenizer = AutoTokenizer.from_pretrained('facebook/bart-large-cnn')
+generator = AutoModelForSeq2SeqLM.from_pretrained('facebook/bart-large-cnn')
+
+class Query(BaseModel):
+    question: str
+    documents: list[str]
+
+@app.post("/generate")
+async def generate_answer(query: Query):
+    if not query.documents:
+        raise HTTPException(status_code=400, detail="No documents provided.")
+
+    # Encode the documents and the query
+    doc_embeddings = retriever.encode(query.documents, convert_to_tensor=True)
+    query_embedding = retriever.encode(query.question, convert_to_tensor=True)
+
+    # Compute cosine similarities
+    similarities = util.pytorch_cos_sim(query_embedding, doc_embeddings)[0]
+    top_doc_index = torch.argmax(similarities).item()
+    top_doc = query.documents[top_doc_index]
+
+    # Prepare input for the generator
+    input_text = f"question: {query.question} context: {top_doc}"
+    inputs = tokenizer(input_text, return_tensors="pt", max_length=512, truncation=True)
+
+    # Generate the answer
+    output_ids = generator.generate(inputs.input_ids, max_length=150, num_beams=5, early_stopping=True)
+    answer = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+
+    return {"answer": answer}

requirements.txt

@@ -0,0 +1,6 @@
+fastapi==0.99.1
+uvicorn[standard]==0.22.0
+transformers==4.33.3
+sentence-transformers==2.2.2
+torch==2.0.1
+requests==2.31.0