app.py

from fastapi import FastAPI, Request
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
import torch
import uvicorn
import os

# Create the FastAPI instance
app = FastAPI(title="Biomarker Apollo API")

# Allow cross-origin requests
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Restrict to specific domains in production
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Model loading
model_loaded = False
try:
    # Check if we're running locally or on HuggingFace Spaces
    if os.path.exists("./adapter_model.bin") and os.path.exists("./tokenizer.model"):
        # Local path
        model_path = "./"
    else:
        # HuggingFace Hub path
        model_path = "opium80s/biomarker-apollo"
    
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    model = AutoModelForCausalLM.from_pretrained(model_path)
    generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
    model_loaded = True
    print("Model loaded successfully!")
except Exception as e:
    print(f"Error loading model: {e}")
    generator = None

class GenerationRequest(BaseModel):
    prompt: str
    max_length: int = 512
    temperature: float = 0.7
    num_return_sequences: int = 1

@app.get("/")
async def root():
    return {"message": "Biomarker Apollo API", "model_loaded": model_loaded}

@app.post("/generate/")
async def generate_text(request: GenerationRequest):
    if not model_loaded:
        return {"error": "Model not loaded successfully"}
    
    try:
        outputs = generator(
            request.prompt,
            max_length=request.max_length,
            temperature=request.temperature,
            num_return_sequences=request.num_return_sequences,
            do_sample=True
        )
        
        # Clean up the output
        generated_texts = [output["generated_text"] for output in outputs]
        
        return {"generated_texts": generated_texts}
    
    except Exception as e:
        return {"error": str(e)}

@app.get("/health")
async def health_check():
    return {"status": "healthy", "model_loaded": model_loaded}

# For local testing only
if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)