imurra's picture
unzip database
a9d88a0 verified
raw
history blame
1.3 kB
# app.py - Hugging Face Spaces version with auto-extract
import os
import zipfile
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import chromadb
from sentence_transformers import SentenceTransformer
import gradio as gr
# Database path
DB_PATH = "./medqa_db"
ZIP_PATH = "./medqa_db.zip"
# Extract database if needed
if not os.path.exists(DB_PATH) and os.path.exists(ZIP_PATH):
print("Extracting database from zip file...")
with zipfile.ZipFile(ZIP_PATH, 'r') as zip_ref:
zip_ref.extractall(".")
print("Database extracted successfully!")
# Initialize
print(f"Loading database from: {DB_PATH}")
client = chromadb.PersistentClient(path=DB_PATH)
collection = client.get_collection("medqa")
print(f"Collection loaded with {collection.count()} items")
print(f"Loading MedCPT model...")
model = SentenceTransformer('ncbi/MedCPT-Query-Encoder')
print("Initialization complete!")
# FastAPI app
app = FastAPI(title="MedQA Search API")
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
class SearchRequest(BaseModel):
query: str
num_results: int = 3
class SearchResponse(BaseModel):
results: list