app.py

from fastapi import FastAPI, Query
from transformers import CLIPModel, CLIPProcessor
import torch

# Initialize FastAPI
app = FastAPI()

# Load CLIP model and processor from Hugging Face
model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")

# Load and process document
with open("test.txt", "r", encoding="utf-8") as f:
    sentences = [line.strip() for line in f if line.strip()]

# Encode document sentences
with torch.no_grad():
    sentence_inputs = processor(text=sentences, return_tensors="pt", padding=True, truncation=True)
    sentence_embeddings = model.get_text_features(**sentence_inputs)

@app.get("/")
def welcome():
    return {"message": "CLIP-based Document Retrieval Service is Running!"}

@app.get("/search")
def search(text: str = Query(..., description="Enter your query"), top_k: int = 5):
    with torch.no_grad():
        query_inputs = processor(text=[text], return_tensors="pt", padding=True, truncation=True)
        query_embedding = model.get_text_features(**query_inputs)

        # Compute cosine similarity
        scores = torch.nn.functional.cosine_similarity(query_embedding, sentence_embeddings)[0]
        top_indices = torch.topk(scores, k=top_k).indices

        results = [{
            "matched_sentence": sentences[i],
            "similarity_score": round(scores[i].item(), 3)
        } for i in top_indices]

    return {
        "query": text,
        "top_matches": results
    }