from fastapi import FastAPI, Form

app = FastAPI()

from llama_cpp import Llama
import time
import os

llm = Llama(model_path="llama-2-7b-chat.Q3_K_S.gguf", n_ctx=2048, n_batch=512, use_mlock=True, n_threads=8)
time.sleep(8)
os.system("cls")
print("Chatbot by Aritra Roy & DVLH")
import warnings
warnings.filterwarnings("ignore")

@app.post("/chat")
async def chat_endpoint(ask: str = Form(...)):
    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)

    try:
        while True:
            try:
                chunk = next(output_stream)
                if chunk.get('choices') and chunk['choices'][0].get('text'):
                    response_text_chunk = chunk['choices'][0]['text']
                    print(response_text_chunk, delay=0.00)  # live response
                    return {"response": response_text_chunk}
            except StopIteration:
                break
    except StopIteration:
        pass

if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="127.0.0.1", port=8000)