Rename app.py to main.py

app.py

@@ -1,79 +0,0 @@
-import gradio as gr
-from typing import List
-import requests
-import json
-import subprocess
-from multiprocessing import Process
-
-def get_text_embedding(text: str, model: str = "mxbai-embed-large", api_url: str = "http://localhost:11434/api/embeddings") -> List[float]:
-    """
-    Sends a prompt to the embedding API and retrieves the embedding.
-
-    Args:
-        text (str): The text to embed.
-        model (str): The model to use for generating the embedding (default is "mxbai-embed-large").
-        api_url (str): The API endpoint URL (default is "http://localhost:11434/api/embeddings").
-
-    Returns:
-        list: A list of floats representing the embedding vector.
-
-    Raises:
-        Exception: If the API request fails.
-    """
-    payload = {
-        "model": model,
-        "prompt": text
-    }
-
-    try:
-        response = requests.post(api_url, data=json.dumps(payload), headers={"Content-Type": "application/json"})
-        response.raise_for_status()  # Raise an error for non-200 status codes
-        data = response.json()
-        return data.get("embedding", [])
-    except requests.exceptions.RequestException as e:
-        raise Exception(f"Error communicating with the embedding API: {e}")
-
-def process_text_to_embedding(text: str) -> str:
-    """Process the text input and return the embedding as a string."""
-    try:
-        embedding = get_text_embedding(text)
-        return json.dumps(embedding, indent=2)
-    except Exception as e:
-        return f"Error: {str(e)}"
-
-def run_ollama_serve():
-    subprocess.run(["ollama", "serve"], check=True)
-
-# Create processes
-serve_process = Process(target=run_ollama_serve)
-
-# Start processes
-serve_process.start()
-
-# subprocess.run(["sudo", "apt", "install", "-y", "pciutils", "lshw"], check=True)
-# subprocess.run(["curl", "-fsSL", "https://ollama.com/install.sh", "|", "sh"], shell=True, check=True)
-subprocess.run(["ollama", "pull", "snowflake-arctic-embed2"], check=True)
-
-# Define the Gradio interface
-def main():
-    title = "Text Embedding Generator"
-    description = "Enter a text input, and this tool will generate an embedding using the specified model via API."
-
-    with gr.Blocks() as demo:
-        gr.Markdown(f"# {title}")
-        gr.Markdown(description)
-
-        with gr.Row():
-            text_input = gr.Textbox(label="Input Text", placeholder="Enter your text here")
-        
-        with gr.Row():
-            output = gr.Textbox(label="Embedding Output", lines=10)
-
-        submit_button = gr.Button("Generate Embedding")
-
-        submit_button.click(fn=process_text_to_embedding, inputs=[text_input], outputs=[output])
-
-    demo.launch()
-
-if __name__ == "__main__":
-    main()

main.py

@@ -0,0 +1,41 @@
+# main.py
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+from typing import List
+from vllm import LLM
+import numpy as np
+
+# Initialize the model
+llm = LLM(model='BAAI/bge-base-en-v1.5', task="embed")
+
+# Initialize FastAPI app
+app = FastAPI()
+
+# Define request schemas
+class DocumentsRequest(BaseModel):
+    documents: List[str]
+
+class QueryRequest(BaseModel):
+    query: str
+
+# API to embed documents
+@app.post("/embed_documents")
+def embed_documents(request: DocumentsRequest):
+    try:
+        docs = request.documents
+        docs_embd = llm.encode(docs)
+        docs_embd = [doc.outputs.data.numpy().tolist() for doc in docs_embd]
+        return {"embeddings": docs_embd}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Error embedding documents: {str(e)}")
+
+# API to embed query
+@app.post("/embed_query")
+def embed_query(request: QueryRequest):
+    try:
+        query = request.query
+        query_embd = llm.encode(query)
+        query_embd = query_embd[0].outputs.data.numpy().tolist()
+        return {"embedding": query_embd}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Error embedding query: {str(e)}")