Update main.py

main.py

@@ -1,105 +1,41 @@
-from torch import Tensor
-from transformers import AutoTokenizer, AutoModel
-from ctranslate2 import Translator
 from typing import Union
-
 from fastapi import FastAPI
 from pydantic import BaseModel
-
-
-def average_pool(last_hidden_states: Tensor,
-                 attention_mask: Tensor) -> Tensor:
-    last_hidden = last_hidden_states.masked_fill(
-        ~attention_mask[..., None].bool(), 0.0)
-    return last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]
-
-
-# text-ada replacement
-embeddingTokenizer = AutoTokenizer.from_pretrained(
-    './multilingual-e5-base')
-embeddingModel = AutoModel.from_pretrained('./multilingual-e5-base')
-
-# chatGpt replacement
-inferenceTokenizer = AutoTokenizer.from_pretrained(
-    "./flan-alpaca-gpt4-xl-ct2")
-inferenceTranslator = Translator(
-    "./flan-alpaca-gpt4-xl-ct2", compute_type="int8", device="cpu")
-
-
-class EmbeddingRequest(BaseModel):
-    input: Union[str, None] = None
-
-
-class TokensCountRequest(BaseModel):
-    input: Union[str, None] = None
+from llama_cpp import Llama
 
 
 class InferenceRequest(BaseModel):
     input: Union[str, None] = None
-    max_length: Union[int, None] = 0
+    max_tokens: Union[int, None] = 0
 
 
 app = FastAPI()
 
+llm = Llama(model_path="./models/vicuna-7b-v1.5.Q4_K_M.gguf",
+            verbose=False, n_ctx=4096)
+
 
 @app.get("/")
 async def root():
     return {"message": "Hello World"}
 
 
-@app.post("/text-embedding")
-async def text_embedding(request: EmbeddingRequest):
-    input = request.input
-
-    # Process the input data
-    batch_dict = embeddingTokenizer([input], max_length=512,
-                                    padding=True, truncation=True, return_tensors='pt')
-    outputs = embeddingModel(**batch_dict)
-    embeddings = average_pool(outputs.last_hidden_state,
-                              batch_dict['attention_mask'])
-
-    # create response
-    return {
-        'embedding': embeddings[0].tolist()
-    }
-
-
 @app.post('/inference')
 async def inference(request: InferenceRequest):
     input_text = request.input
-    max_length = 256
+    max_tokens = 256
     try:
-        max_length = int(request.max_length)
-        max_length = min(1024, max_length)
+        max_tokens = int(request.max_tokens)
     except:
         pass
 
     # process request
-    input_tokens = inferenceTokenizer.convert_ids_to_tokens(
-        inferenceTokenizer.encode(input_text))
-
-    results = inferenceTranslator.translate_batch(
-        [input_tokens], beam_size=1, max_input_length=0, max_decoding_length=max_length, num_hypotheses=1, repetition_penalty=1.3, sampling_topk=40, sampling_temperature=0.7, use_vmap=False)
-
-    output_tokens = results[0].hypotheses[0]
-    output_text = inferenceTokenizer.decode(
-        inferenceTokenizer.convert_tokens_to_ids(output_tokens), skip_special_tokens=True)
-
-    # create response
-    return {
-        'generated_text': output_text
-    }
-
-
-@app.post('/tokens-count')
-async def tokens_count(request: TokensCountRequest):
-    input_text = request.input
-
-    tokens = inferenceTokenizer.convert_ids_to_tokens(
-        inferenceTokenizer.encode(input_text))
+    try:
+        result = llm(input_text, temperature=0.2,
+                     top_k=5, max_tokens=max_tokens)
+        return result
+    except:
+        pass
 
     # create response
-    return {
-        'tokens': tokens,
-        'total': len(tokens)
-    }
+    return {}