Update app.py

app.py

@@ -1,39 +1,29 @@
 from fastapi import FastAPI, HTTPException
 from transformers import AutoModelForCausalLM, AutoTokenizer
-from huggingface_hub import hf_hub_download
-from huggingface_hub import snapshot_download
 import os
+import copy
+import time
+import llama_cpp
+from llama_cpp import Llama
+from huggingface_hub import hf_hub_download  
 
 app = FastAPI()
 
+llm = Llama(
+    model_path=hf_hub_download(
+        repo_id="TheBloke/Llama-2-7b-Chat-GGUF",
+        filename="mistral-7b-v0.1.Q4_K_M.gguf"),
+    n_ctx=2048,
+) 
+
 @app.get("/")
 async def generate_text():
     try:
-        # Get the cache directory from the environment variable
-        cache_dir = os.getenv("HF_HOME")
-
-        # Specify the directory for model download within the Docker container
-        model_dir = os.path.join(cache_dir, "TheBloke/Mistral-7B-v0.1-GGUF")
-        os.makedirs(model_dir, exist_ok=True)
-        
-        #hf_hub_download(repo_id="TheBloke/Mistral-7B-v0.1-GGUF", filename="mistral-7b-v0.1.Q4_K_M.gguf", local_dir=model_dir)
-        #hf_hub_download(repo_id="TheBloke/Mistral-7B-v0.1-GGUF", filename="config.json", local_dir=model_dir)
-        #snapshot_download(repo_id="TheBloke/Mistral-7B-v0.1-GGUF", local_dir=model_dir)
-
-        # Check if config.json file exists in the model directory
-        config_file = os.path.join(model_dir, "config.json")
-        if not os.path.exists(config_file):
-            raise ValueError("config.json file is missing in the model directory")
-
-        # Load tokenizer and model
-        tokenizer = AutoTokenizer.from_pretrained(model_dir)
-        model = AutoModelForCausalLM.from_pretrained(model_dir)
-
-        # Generate text
-        prompt = "Once upon a time, there was a"
-        inputs = tokenizer(prompt, return_tensors="pt")
-        output = model.generate(input_ids=inputs["input_ids"], max_length=50, num_return_sequences=3, temperature=0.7)
-        generated_texts = tokenizer.batch_decode(output, skip_special_tokens=True)
-        return generated_texts
+        output = llm(
+          "Q: Name the planets in the solar system? A: ", 
+          max_tokens=32, 
+          stop=["Q:", "\n"], 
+          echo=True)
+        return output
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))