Tonyivan commited on
Commit
bbd40ae
·
verified ·
1 Parent(s): f57466f

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +11 -15
app.py CHANGED
@@ -1,4 +1,4 @@
1
- from fastapi import FastAPI, HTTPException
2
  from pydantic import BaseModel
3
  from sentence_transformers import SentenceTransformer, util
4
  from transformers import pipeline
@@ -14,9 +14,6 @@ nlp = pipeline('question-answering', model=question_model, tokenizer=question_mo
14
 
15
  summarizer = pipeline("summarization", model="facebook/bart-large-cnn")
16
 
17
- # Request models
18
- class ModifyQueryRequest(BaseModel):
19
- query_string: str
20
 
21
  class ModifyQueryRequest_v3(BaseModel):
22
  query_string_list: List[str]
@@ -29,9 +26,6 @@ class AnswerQuestionRequest(BaseModel):
29
  class T5QuestionRequest(BaseModel):
30
  context: str
31
 
32
- # Response models
33
- class ModifyQueryResponse(BaseModel):
34
- embeddings: List[List[float]]
35
 
36
  class AnswerQuestionResponse(BaseModel):
37
  answer: str
@@ -41,20 +35,22 @@ class T5Response(BaseModel):
41
  answer: str
42
 
43
  # API endpoints
44
- @app.post("/modify_query", response_model=ModifyQueryResponse)
45
- async def modify_query(request: ModifyQueryRequest):
46
  try:
47
- binary_embeddings = model.encode([request.query_string], precision="binary")
48
- return ModifyQueryResponse(embeddings=binary_embeddings[0].tolist())
 
49
  except Exception as e:
50
  raise HTTPException(status_code=500, detail=str(e))
51
 
52
- @app.post("/modify_query_v3", response_model=ModifyQueryResponse)
53
- async def modify_query_v3(request: ModifyQueryRequest_v3):
54
  try:
55
  # Generate embeddings for a list of query strings
56
- embeddings = model.encode(request.query_string_list)
57
- return ModifyQueryResponse(embeddings=[emb.tolist() for emb in embeddings])
 
58
  except Exception as e:
59
  raise HTTPException(status_code=500, detail=f"Error in modifying query v3: {str(e)}")
60
 
 
1
+ from fastapi import FastAPI, HTTPException, JSONResponse, Request
2
  from pydantic import BaseModel
3
  from sentence_transformers import SentenceTransformer, util
4
  from transformers import pipeline
 
14
 
15
  summarizer = pipeline("summarization", model="facebook/bart-large-cnn")
16
 
 
 
 
17
 
18
  class ModifyQueryRequest_v3(BaseModel):
19
  query_string_list: List[str]
 
26
  class T5QuestionRequest(BaseModel):
27
  context: str
28
 
 
 
 
29
 
30
  class AnswerQuestionResponse(BaseModel):
31
  answer: str
 
35
  answer: str
36
 
37
  # API endpoints
38
+ @app.post("/modify_query")
39
+ async def modify_query(request: Request):
40
  try:
41
+ raw_data = await request.json()
42
+ binary_embeddings = model.encode([raw_data['query_string']], precision="binary")
43
+ return JSONResponse(content={'embeddings':binary_embeddings[0].tolist()})
44
  except Exception as e:
45
  raise HTTPException(status_code=500, detail=str(e))
46
 
47
+ @app.post("/modify_query_v3")
48
+ async def modify_query_v3(request: Request):
49
  try:
50
  # Generate embeddings for a list of query strings
51
+ raw_data = await request.json()
52
+ embeddings = model.encode(raw_data['query_string_list'])
53
+ return JSONResponse(content={'embeddings':[emb.tolist() for emb in embeddings]})
54
  except Exception as e:
55
  raise HTTPException(status_code=500, detail=f"Error in modifying query v3: {str(e)}")
56