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))