khalednabawi11 commited on
Commit
c395c47
·
verified ·
1 Parent(s): 1412372

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +4 -1
app.py CHANGED
@@ -268,6 +268,8 @@ generation_config = GenerationConfig(
268
  pipe = TextGenerationPipeline(
269
  model=model,
270
  tokenizer=tokenizer,
 
 
271
  device=model.device.index if torch.cuda.is_available() else -1
272
  )
273
 
@@ -344,7 +346,8 @@ async def chat_fn(query: Query):
344
  @app.post("/ask-rag")
345
  async def chat_fn(query: Query):
346
  message = query.message
347
- logger.info(f"Received message: {message}")
 
348
 
349
  # Run RAG inference in thread
350
  loop = asyncio.get_event_loop()
 
268
  pipe = TextGenerationPipeline(
269
  model=model,
270
  tokenizer=tokenizer,
271
+ generation_config = generation_config,
272
+ task = "text-generation",
273
  device=model.device.index if torch.cuda.is_available() else -1
274
  )
275
 
 
346
  @app.post("/ask-rag")
347
  async def chat_fn(query: Query):
348
  message = query.message
349
+ prompt = generate_prompt(message)
350
+ logger.info(f"Received message: {prompt}")
351
 
352
  # Run RAG inference in thread
353
  loop = asyncio.get_event_loop()