from fastapi import FastAPI
from pydantic import BaseModel
import torch
from transformers import AutoTokenizer, AutoModel
import uvicorn
import gradio as gr

# Initialize FastAPI app
app = FastAPI()

# Load pre-trained model and tokenizer
model_name = "bert-base-uncased"  # You can change this to another model

try:
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModel.from_pretrained(model_name)
except Exception as e:
    print(f"Error loading model: {e}")

class TextRequest(BaseModel):
    text: str

# Function to generate embeddings
def get_embeddings(text: str):
    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
    with torch.no_grad():
        outputs = model(**inputs)
    embeddings = outputs.last_hidden_state.mean(dim=1)  # Pooling
    return embeddings.numpy().tolist()  # Convert to list for API response

@app.post("/get-embedding/")
async def get_embedding(request: TextRequest):
    text = request.text
    embeddings = get_embeddings(text)
    return {"embedding": embeddings}

def gradio_interface(text):
    return get_embeddings(text)

grn = gr.Interface(fn=gradio_interface, inputs="text", outputs="json", title="Text Embedding Generator")

grn.launch(server_name="0.0.0.0", server_port=7860, share=True)

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)  # Port 7860 for Hugging Face Spaces