app

app.py

@@ -1,5 +1,6 @@
 from fastapi import FastAPI
 from transformers import AutoModelForCausalLM, AutoTokenizer
+from llama_cpp import Llama
 
 
 # Create an instance of the FastAPI class
@@ -8,18 +9,16 @@ app = FastAPI()
 # Define a route for the root endpoint
 @app.get("/llm")
 async def read_root():
-    device = "cpu"
-    model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-Instruct-v0.1")
-    tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.1")
-    text = """<s>[INST] What is your favourite condiment? [/INST]
-    """
-    encodeds = tokenizer(text, return_tensors="pt", add_special_tokens=False)
-    model_inputs = encodeds.to(device)
-    model.to(device)
-    generated_ids = model.generate(**model_inputs, max_new_tokens=1000, do_sample=True)
-    decoded = tokenizer.batch_decode(generated_ids)
-    print(decoded[0])
-    return {"message": decoded[0]}
+    llm = Llama.from_pretrained(
+    repo_id="Qwen/Qwen1.5-0.5B-Chat-GGUF",
+    filename="*q8_0.gguf",
+    verbose=False)
+    output = llm(
+      "Q: Name the planets in the solar system? A: ", # Prompt
+      max_tokens=32, # Generate up to 32 tokens, set to None to generate up to the end of the context window
+      stop=["Q:", "\n"], # Stop generating just before the model would generate a new question
+      echo=True # Echo the prompt back in the output)
+    return {"message": output}