ka1kuk commited on
Commit
2f57100
1 Parent(s): 6f97008

Update apis/chat_api.py

Browse files
Files changed (1) hide show
  1. apis/chat_api.py +4 -4
apis/chat_api.py CHANGED
@@ -125,9 +125,9 @@ class ChatAPIApp:
125
  return None
126
 
127
  class QueryRequest(BaseModel):
128
- texts: List[str]
129
  model_name: str = Field(..., example="bert-base-uncased")
130
- api_key: str = Field(..., example="your_hf_api_key_here")
131
 
132
  class ChatCompletionsPostItem(BaseModel):
133
  model: str = Field(
@@ -187,7 +187,7 @@ class ChatAPIApp:
187
  data_response = streamer.chat_return_dict(stream_response)
188
  return data_response
189
 
190
- async def chat_embedding(self, texts, model_name, api_key):
191
  api_url = f"https://api-inference.huggingface.co/pipeline/feature-extraction/{model_name}"
192
  headers = {"Authorization": f"Bearer {api_key}"}
193
  response = requests.post(api_url, headers=headers, json={"inputs": texts})
@@ -206,7 +206,7 @@ class ChatAPIApp:
206
  try:
207
  for attempt in range(3): # Retry logic
208
  try:
209
- embeddings = await self.chat_embedding(request.texts, request.model_name, request.api_key)
210
  data = [
211
  {"object": "embedding", "index": i, "embedding": embedding}
212
  for i, embedding in enumerate(embeddings)
 
125
  return None
126
 
127
  class QueryRequest(BaseModel):
128
+ input: str
129
  model_name: str = Field(..., example="bert-base-uncased")
130
+ encoding_format: str
131
 
132
  class ChatCompletionsPostItem(BaseModel):
133
  model: str = Field(
 
187
  data_response = streamer.chat_return_dict(stream_response)
188
  return data_response
189
 
190
+ async def chat_embedding(self, texts, model_name, api_key, api_key: str = Depends(extract_api_key)):
191
  api_url = f"https://api-inference.huggingface.co/pipeline/feature-extraction/{model_name}"
192
  headers = {"Authorization": f"Bearer {api_key}"}
193
  response = requests.post(api_url, headers=headers, json={"inputs": texts})
 
206
  try:
207
  for attempt in range(3): # Retry logic
208
  try:
209
+ embeddings = await self.chat_embedding(request.input, request.model, request.api_key)
210
  data = [
211
  {"object": "embedding", "index": i, "embedding": embedding}
212
  for i, embedding in enumerate(embeddings)