from fastapi import FastAPI
from transformers import pipeline
import os
from huggingface_hub import login

# Load the Hugging Face token from environment variables
hf_token = os.getenv("HF_TOKEN")
if hf_token:
    login(token=hf_token)

# Initialize the Hugging Face pipeline
pipe = pipeline("text-generation", model="microsoft/phi-4", trust_remote_code=True)

# FastAPI app setup
app = FastAPI()

@app.get('/')
def home():
    return {"hello": "Geominds"}

@app.get('/ask')
def ask(prompt: str):
    result = pipe(prompt)
    return result[0]