Update app.py

app.py

@@ -1,23 +1,25 @@
-from fastapi import FastAPI
+from fastapi import FastAPI, HTTPException
 from transformers import AutoModelForCausalLM, AutoTokenizer
-from llama_cpp import Llama
-from huggingface_hub import hf_hub_download
+import os
 
 app = FastAPI()
 
-@app.get("/")
-async def install_model():
-    #hf_hub_download(repo_id="TheBloke/Mistral-7B-v0.1-GGUF", filename="mistral-7b-v0.1.Q4_K_M.gguf")
-    prompt = "Once upon a time, there was a"
-    tokenizer = AutoTokenizer.from_pretrained("./mistral-7b-v0.1.Q4_K_M.gguf")
-    model = AutoModelForCausalLM.from_pretrained("./mistral-7b-v0.1.Q4_K_M.gguf")
-    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)
-    for i, text in enumerate(generated_texts):
-        print(f"Generated Text {i+1}: {text}")
-    return generated_texts
-
-
+cache_dir = os.path.expanduser("~/.cache/huggingface/")
+model_folder = os.path.join(cache_dir, "TheBloke/Mistral-7B-v0.1-GGUF")
+if not os.path.exists(model_folder):
+    raise ValueError("Model not found in cache directory. Please download the model.")
 
+# Load the tokenizer and model asynchronously
+tokenizer = AutoTokenizer.from_pretrained(model_folder)
+model = AutoModelForCausalLM.from_pretrained(model_folder)
 
+@app.get("/")
+async def generate_text():
+    try:
+        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
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))