Update app.py

app.py

@@ -260,45 +260,22 @@ def simple_chat(message: dict, temperature: float = 0.8, max_length: int = 4096,
         gen_kwargs = {**input_ids, **generate_kwargs}
 
         with torch.no_grad():
-            thread = Thread(target=model.generate, kwargs=gen_kwargs)
-            thread.start()
-            buffer = ""
-            for new_text in streamer:
-                buffer += new_text
+            generated_ids = model.generate(input_ids['input_ids'], **generate_kwargs)
+            generated_text = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
 
-        return buffer
 
+        print("---------")
+        print(generated_ids[0])
+        print("---------")
+        print(generated_text)
+        print("---------")
+
+
+        return PlainTextResponse(generated_text)
     except Exception as e:
-        return f"Error: {str(e)}"
-
-@app.post("/chat/")
-async def test_endpoint(
-    text: str = Form(...), 
-    file: UploadFile = File(None)
-):
-    if file:
-        file_content = BytesIO(await file.read())
-        file_name = file.filename
-        
-        message = {
-            "text": text,
-            "file_content": file_content,
-            "file_name": file_name
-        }
-    else:
-        message = {
-            "text": text,
-            "file_content": None,
-            "file_name": None
-        }
-    
-    print(message)
-    response = simple_chat(message)
-    
-    if isinstance(response, str) and response.startswith("Error:"):
-        return PlainTextResponse(response, status_code=500)
-    else:
-        return StreamingResponse(BytesIO(response.encode()), media_type="text/plain")
+        return PlainTextResponse(f"Error: {str(e)}")
+
+
         
 # def simple_chat(message: dict, temperature: float = 0.8, max_length: int = 4096, top_p: float = 1, top_k: int = 10, penalty: float = 1.0):
 #     try: