Update app.py

app.py

@@ -11,7 +11,7 @@ import asyncio
 import gradio as gr
 import torch
 from dotenv import load_dotenv
-from fastapi import FastAPI
+from fastapi import FastAPI, Request
 from fastapi.responses import JSONResponse
 import uvicorn
 from threading import Thread
@@ -204,14 +204,10 @@ def get_best_response(responses):
     best_response_index = total_similarities.argmax()
     return responses[best_response_index]
 
-async def generate_model_response(model, inputs, max_tokens=2048):
+async def generate_model_response(model, inputs):
     try:
         response = model(inputs)
-        text = remove_duplicates(response['choices'][0]['text'])
-        
-        if len(text.split()) > max_tokens:
-            return text[:max_tokens]
-        return text
+        return remove_duplicates(response['choices'][0]['text'])
     except Exception as e:
         return ""
 
@@ -232,9 +228,10 @@ async def process_message(message):
 app = FastAPI()
 
 @app.post("/generate")
-async def generate(request: BaseModel):
+async def generate(request: Request):
     try:
-        response = await process_message(request.message)
+        body = await request.json()
+        response = await process_message(body['message'])
         return JSONResponse(content={"response": response})
     except Exception as e:
         return JSONResponse(content={"error": str(e)})