Update main.py

main.py

@@ -1,38 +1,26 @@
+
+from typing import Union
+from wrapper import LLMWrapper
 import uvicorn
 from fastapi import FastAPI, Request
-from langchain_community.llms.ctransformers import CTransformers
-
-
-
-MODEL_TYPE = 'mistral'
-MODEL_BIN_PATH = "mistral-7b-instruct-v0.1.Q3_K_S.gguf"
-MAX_NEW_TOKEN =  600
-TEMPRATURE =  0.01
-CONTEXT_LENGTH =  6000
-
 
-llm = CTransformers(
-            model=MODEL_BIN_PATH,
-            config={
-                'max_new_tokens': MAX_NEW_TOKEN,
-                'temperature': TEMPRATURE,
-                'context_length': CONTEXT_LENGTH
-            },
-            model_type=MODEL_TYPE
-        )
 
 app = FastAPI()
+llm_wrapper = LLMWrapper()
 
 
-@app.post("//llm_on_cpu")
+
+@app.post("/")
 async def generate_text(request: Request):
     raw_data = await request.body()  # Get the raw body data from the request
     prompt = raw_data.decode('utf-8')
     if not prompt:
         return {'error': 'Prompt is required'}, 400
-    try:
-        generated_text = llm(prompt)
-        print(generated_text)
-        return {'generated_text': generated_text}
-    except:
-        return {"generated_text": "Error!"}
+
+    generated_text = llm_wrapper.generate_text(prompt)
+    print(generated_text)
+    return {'generated_text': generated_text}
+
+
+if __name__ == '__main__':
+    uvicorn.run(app, host='127.0.0.1', port=8001)