Tonyivan commited on
Commit
9b74ec6
1 Parent(s): 9dbbc34

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +52 -0
app.py ADDED
@@ -0,0 +1,52 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from fastapi import FastAPI, HTTPException
2
+ from pydantic import BaseModel
3
+ from sentence_transformers import SentenceTransformer
4
+ from transformers import pipeline
5
+
6
+ # Initialize FastAPI app
7
+ app = FastAPI()
8
+
9
+ # Load models
10
+ model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
11
+ question_model = "deepset/tinyroberta-squad2"
12
+ nlp = pipeline('question-answering', model=question_model, tokenizer=question_model)
13
+
14
+ # Define request models
15
+ class ModifyQueryRequest(BaseModel):
16
+ query_string: str
17
+
18
+ class AnswerQuestionRequest(BaseModel):
19
+ question: str
20
+ context: str
21
+
22
+ # Define response models (if needed)
23
+ class ModifyQueryResponse(BaseModel):
24
+ embeddings: list
25
+
26
+ class AnswerQuestionResponse(BaseModel):
27
+ answer: str
28
+
29
+ # Define API endpoints
30
+ @app.post("/modify_query", response_model=ModifyQueryResponse)
31
+ async def modify_query(request: ModifyQueryRequest):
32
+ try:
33
+ binary_embeddings = model.encode([request.query_string], precision="binary")
34
+ return ModifyQueryResponse(embeddings=binary_embeddings[0].tolist())
35
+ except Exception as e:
36
+ raise HTTPException(status_code=500, detail=str(e))
37
+
38
+ @app.post("/answer_question", response_model=AnswerQuestionResponse)
39
+ async def answer_question(request: AnswerQuestionRequest):
40
+ try:
41
+ QA_input = {
42
+ 'question': request.question,
43
+ 'context': request.context
44
+ }
45
+ result = nlp(QA_input)
46
+ return AnswerQuestionResponse(answer=result['answer'])
47
+ except Exception as e:
48
+ raise HTTPException(status_code=500, detail=str(e))
49
+
50
+ if __name__ == "__main__":
51
+ import uvicorn
52
+ uvicorn.run(app, host="0.0.0.0", port=8000)