fashxp commited on
Commit
afab4d9
·
1 Parent(s): 7bac21a

made tailing / optional

Browse files
Files changed (1) hide show
  1. src/main.py +12 -6
src/main.py CHANGED
@@ -63,7 +63,7 @@ async def gpu_check():
63
  # Translation Task
64
  # =========================
65
  @app.post(
66
- "/translation/{model_name:path}/",
67
  openapi_extra={
68
  "requestBody": {
69
  "content": {
@@ -92,6 +92,7 @@ async def translate(
92
  list: The translation result(s) as returned by the pipeline.
93
  """
94
 
 
95
  translationTaskService = TranslationTaskService(logger)
96
  return await translationTaskService.translate(request, model_name)
97
 
@@ -100,7 +101,7 @@ async def translate(
100
  # Zero-Shot Image Classification Task
101
  # =========================
102
  @app.post(
103
- "/zero-shot-image-classification/{model_name:path}/",
104
  openapi_extra={
105
  "requestBody": {
106
  "content": {
@@ -129,6 +130,7 @@ async def zero_shot_image_classification(
129
  list: The classification result(s) as returned by the pipeline.
130
  """
131
 
 
132
  zeroShotTask = ClassificationTaskService(logger, 'zero-shot-image-classification')
133
  return await zeroShotTask.classify(request, model_name)
134
 
@@ -137,7 +139,7 @@ async def zero_shot_image_classification(
137
  # Image Classification Task
138
  # =========================
139
  @app.post(
140
- "/image-classification/{model_name:path}/",
141
  openapi_extra={
142
  "requestBody": {
143
  "content": {
@@ -165,6 +167,7 @@ async def image_classification(
165
  list: The classification result(s) as returned by the pipeline.
166
  """
167
 
 
168
  imageTask = ClassificationTaskService(logger, 'image-classification')
169
  return await imageTask.classify(request, model_name)
170
 
@@ -174,7 +177,7 @@ async def image_classification(
174
  # Zero-Shot Text Classification Task
175
  # =========================
176
  @app.post(
177
- "/zero-shot-text-classification/{model_name:path}/",
178
  openapi_extra={
179
  "requestBody": {
180
  "content": {
@@ -203,6 +206,7 @@ async def zero_shot_text_classification(
203
  list: The classification result(s) as returned by the pipeline.
204
  """
205
 
 
206
  zeroShotTask = ClassificationTaskService(logger, 'zero-shot-classification')
207
  return await zeroShotTask.classify(request, model_name)
208
 
@@ -211,7 +215,7 @@ async def zero_shot_text_classification(
211
  # Text Classification Task
212
  # =========================
213
  @app.post(
214
- "/text-classification/{model_name:path}/",
215
  openapi_extra={
216
  "requestBody": {
217
  "content": {
@@ -239,6 +243,7 @@ async def text_classification(
239
  list: The classification result(s) as returned by the pipeline.
240
  """
241
 
 
242
  textTask = ClassificationTaskService(logger, 'text-classification')
243
  return await textTask.classify(request, model_name)
244
 
@@ -250,7 +255,7 @@ async def text_classification(
250
  # Image to Text Task
251
  # =========================
252
  @app.post(
253
- "/image-to-text/{model_name:path}/",
254
  openapi_extra={
255
  "requestBody": {
256
  "content": {
@@ -286,5 +291,6 @@ async def image_to_text(
286
  list: The generated text as returned by the pipeline.
287
  """
288
 
 
289
  imageToTextTask = TextToImageTaskService(logger)
290
  return await imageToTextTask.extract(request, model_name)
 
63
  # Translation Task
64
  # =========================
65
  @app.post(
66
+ "/translation/{model_name:path}",
67
  openapi_extra={
68
  "requestBody": {
69
  "content": {
 
92
  list: The translation result(s) as returned by the pipeline.
93
  """
94
 
95
+ model_name = model_name.rstrip("/")
96
  translationTaskService = TranslationTaskService(logger)
97
  return await translationTaskService.translate(request, model_name)
98
 
 
101
  # Zero-Shot Image Classification Task
102
  # =========================
103
  @app.post(
104
+ "/zero-shot-image-classification/{model_name:path}",
105
  openapi_extra={
106
  "requestBody": {
107
  "content": {
 
130
  list: The classification result(s) as returned by the pipeline.
131
  """
132
 
133
+ model_name = model_name.rstrip("/")
134
  zeroShotTask = ClassificationTaskService(logger, 'zero-shot-image-classification')
135
  return await zeroShotTask.classify(request, model_name)
136
 
 
139
  # Image Classification Task
140
  # =========================
141
  @app.post(
142
+ "/image-classification/{model_name:path}",
143
  openapi_extra={
144
  "requestBody": {
145
  "content": {
 
167
  list: The classification result(s) as returned by the pipeline.
168
  """
169
 
170
+ model_name = model_name.rstrip("/")
171
  imageTask = ClassificationTaskService(logger, 'image-classification')
172
  return await imageTask.classify(request, model_name)
173
 
 
177
  # Zero-Shot Text Classification Task
178
  # =========================
179
  @app.post(
180
+ "/zero-shot-text-classification/{model_name:path}",
181
  openapi_extra={
182
  "requestBody": {
183
  "content": {
 
206
  list: The classification result(s) as returned by the pipeline.
207
  """
208
 
209
+ model_name = model_name.rstrip("/")
210
  zeroShotTask = ClassificationTaskService(logger, 'zero-shot-classification')
211
  return await zeroShotTask.classify(request, model_name)
212
 
 
215
  # Text Classification Task
216
  # =========================
217
  @app.post(
218
+ "/text-classification/{model_name:path}",
219
  openapi_extra={
220
  "requestBody": {
221
  "content": {
 
243
  list: The classification result(s) as returned by the pipeline.
244
  """
245
 
246
+ model_name = model_name.rstrip("/")
247
  textTask = ClassificationTaskService(logger, 'text-classification')
248
  return await textTask.classify(request, model_name)
249
 
 
255
  # Image to Text Task
256
  # =========================
257
  @app.post(
258
+ "/image-to-text/{model_name:path}",
259
  openapi_extra={
260
  "requestBody": {
261
  "content": {
 
291
  list: The generated text as returned by the pipeline.
292
  """
293
 
294
+ model_name = model_name.rstrip("/")
295
  imageToTextTask = TextToImageTaskService(logger)
296
  return await imageToTextTask.extract(request, model_name)