Spaces:
Runtime error
Runtime error
File size: 1,781 Bytes
0314811 37527e9 291ac95 39504a0 0314811 37527e9 f726f33 87e455f 0314811 84927e5 8f4dd5c 84927e5 ac2091a 84927e5 ac2091a 44ba060 f3e9307 df266e0 097833a ac2091a df266e0 097833a 0314811 84927e5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 |
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
app = FastAPI()
@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
except Exception as e:
raise HTTPException(status_code=500, detail=str(e)) |