import os
from fastapi import FastAPI
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel
import torch

# Set writable cache directory inside the container
os.environ['SENTENCE_TRANSFORMERS_HOME'] = '/app/hf_home'
os.environ['TRANSFORMERS_CACHE'] = '/app/hf_home'

# Ensure the directory exists
os.makedirs(os.environ['TRANSFORMERS_CACHE'], exist_ok=True)

# Define base model and adapter model
base_model_name = "facebook/opt-2.7b"
adapter_name = "mynuddin/chatbot"

# Load base model and tokenizer
tokenizer = AutoTokenizer.from_pretrained(base_model_name)
base_model = AutoModelForCausalLM.from_pretrained(base_model_name, torch_dtype=torch.float16)

# Load PEFT adapter
model = PeftModel.from_pretrained(base_model, adapter_name)
model = model.to("cpu")  # Change to "cuda" if running on GPU
model.eval()

app = FastAPI()

@app.post("/generate")
def generate_text(prompt: str):
    inputs = tokenizer(prompt, return_tensors="pt")
    with torch.no_grad():
        output = model.generate(**inputs, max_length=128)
    generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
    return {"generated_query": generated_text}