Update main.py

main.py

@@ -1,8 +1,27 @@
 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()
 
@@ -15,9 +34,5 @@ async def generate_text(request: Request):
     if not prompt:
         return {'error': 'Prompt is required'}, 400
 
-    generated_text = llm_wrapper.generate_text(prompt)
+    generated_text = llm(prompt)
     return {'generated_text': generated_text}
-
-
-if __name__ == '__main__':
-    uvicorn.run(app, host='127.0.0.1', port=8001)