Update app.py

app.py

@@ -1,4 +1,5 @@
 from fastapi import FastAPI, HTTPException, Request
+from fastapi.responses import StreamingResponse
 from llama_cpp import Llama
 import time
 import os
@@ -7,11 +8,20 @@ app = FastAPI()
 
 llm = Llama(model_path="llama-2-7b-chat.Q3_K_S.gguf", n_ctx=2048, n_batch=512, use_mlock=True, n_threads=8)
 
-def typewrite(text, delay=0.01):
-    for char in text:
-        print(char, end='', flush=True)
-        time.sleep(delay)
-    print(end='', flush=True)  # Print newline to move to the next line
+def generate_output_stream(prompt):
+    output_stream = llm(prompt, max_tokens=1024, echo=False, temperature=0.2, top_p=0.1, stream=True)
+    while True:
+        try:
+            chunk = next(output_stream)
+            if chunk.get('choices') and chunk['choices'][0].get('text'):
+                text = chunk['choices'][0]['text']
+                yield text.encode('utf-8')
+                time.sleep(0.01)  # Adjust the delay if needed
+        except StopIteration:
+            break
+
+def clear_screen():
+    os.system("cls")
 
 @app.post("/chat")
 async def chat(request: Request):
@@ -22,33 +32,18 @@ async def chat(request: Request):
         if not user_input:
             raise HTTPException(status_code=400, detail="Missing 'user_input' field in the request JSON.")
 
-        os.system("cls")
-        print("Chatbot by Aritra Roy & DVLH")
-
         ask = user_input
 
         prompt = f"Llama-2-Chat [INST] <<SYS>>You're an assistant named Tusti. You are Developed by Aritra Roy. Don't share any false information.<</SYS>> {ask} [/INST]"
 
-        output_stream = llm(prompt, max_tokens=1024, echo=False, temperature=0.2, top_p=0.1, stream=True)
-
-        # Remove the response_text variable and directly print the chunks
-        while True:
-            try:
-                chunk = next(output_stream)
-                if chunk.get('choices') and chunk['choices'][0].get('text'):
-                    text = chunk['choices'][0]['text']
-                    typewrite(text, delay=0.00)
-            except StopIteration:
-                break
-
         if ask == 'clear':
-            os.system("cls")
+            clear_screen()
 
-        return {} # Since response_text is no longer used, return an empty dictionary
+        return StreamingResponse(generate_output_stream(prompt), media_type="text/plain")
 
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
 if __name__ == "__main__":
     import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    uvicorn.run(app, host="0.0.0.0", port=7860)