import os
from fastapi import FastAPI
from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel, PeftConfig
from huggingface_hub import login

# Get the Hugging Face token from the environment variable
huggingface_token = os.getenv("HUGGING_FACE_TOKEN")
if huggingface_token is None:
    raise ValueError("HUGGING_FACE_TOKEN environment variable is not set")

# Login to Hugging Face Hub
login(token=huggingface_token)

# Initialize FastAPI app
app = FastAPI()

# Load PEFT model configuration and base model
config = PeftConfig.from_pretrained("frankmorales2020/Mistral-7B-text-to-sql-flash-attention-2-dataeval")
base_model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-Instruct-v0.3", use_auth_token=True)
model = PeftModel.from_pretrained(base_model, "frankmorales2020/Mistral-7B-text-to-sql-flash-attention-2-dataeval")

# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.3", use_auth_token=True)

# Create the pipeline
pipe = pipeline("text2sql", model=model, tokenizer=tokenizer)

@app.get("/")
def home():
    return {"message": "Hello World"}

@app.get("/generate")
def generate(text: str):
    output = pipe(text)
    return {"output": output[0]['generated_text']}