Update app.py

app.py

@@ -1,115 +1,33 @@
-import gradio as gr
+from transformers import AutoTokenizer, AutoModel
+import torch
+import torch.nn.functional as F
 
-def greet(name):
-    return f"Hello {name}! Autonomous AI initialized."
+#Mean Pooling - Take attention mask into account for correct averaging
+def mean_pooling(model_output, attention_mask):
+    token_embeddings = model_output[0] #First element of model_output contains all token embeddings
+    input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
+    return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)
 
-with gr.Blocks() as demo:
-    name = gr.Textbox(label="Your Name")
-    output = gr.Textbox(label="Greeting")
-    name.change(greet, name, output)
 
-demo.launch()
+# Sentences we want sentence embeddings for
+sentences = ['This is an example sentence', 'Each sentence is converted']
 
-from fastapi import FastAPI
-from fastapi.middleware.cors import CORSMiddleware
-from fastapi.responses import HTMLResponse
-from fastapi.staticfiles import StaticFiles
-import gradio as gr
-import numpy as np
-import faiss
-import logging
-import os
-import requests
-import json
-from sentence_transformers import SentenceTransformer
+# Load model from HuggingFace Hub
+tokenizer = AutoTokenizer.from_pretrained('sentence-transformers/all-MiniLM-L6-v2')
+model = AutoModel.from_pretrained('sentence-transformers/all-MiniLM-L6-v2')
 
-# === Environment variables (safe cache paths for HF Spaces) ===
-os.environ["TRANSFORMERS_CACHE"] = "/data/transformers"
-os.environ["HF_HOME"] = "/data/hf"
-os.environ["SENTENCE_TRANSFORMERS_HOME"] = "/data/st"
+# Tokenize sentences
+encoded_input = tokenizer(sentences, padding=True, truncation=True, return_tensors='pt')
 
-# === Load embedding model ===
-model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
-index = faiss.IndexFlatL2(384)
-memory_text = []
+# Compute token embeddings
+with torch.no_grad():
+    model_output = model(**encoded_input)
 
-# === Fireworks API ===
-FIREWORKS_API_KEY = os.getenv("FIREWORKS_API_KEY")  # 🔐 Use env var instead of hardcoding
-FIREWORKS_URL = "https://api.fireworks.ai/inference/v1/chat/completions"
+# Perform pooling
+sentence_embeddings = mean_pooling(model_output, encoded_input['attention_mask'])
 
-def query_fireworks(prompt):
-    payload = {
-        "model": "accounts/fireworks/models/deepseek-r1",
-        "max_tokens": 4096,
-        "top_p": 1,
-        "top_k": 40,
-        "temperature": 0.6,
-        "messages": [{"role": "user", "content": prompt}],
-    }
-    headers = {
-        "Accept": "application/json",
-        "Content-Type": "application/json",
-        "Authorization": f"Bearer {FIREWORKS_API_KEY}"
-    }
+# Normalize embeddings
+sentence_embeddings = F.normalize(sentence_embeddings, p=2, dim=1)
 
-    response = requests.post(FIREWORKS_URL, headers=headers, data=json.dumps(payload))
-    result = response.json()
-    return result.get("choices", [{}])[0].get("message", {}).get("content", "⚠️ No response.")
-
-# === Autonomous Agent ===
-def autonomous_agent(input_text):
-    vec = model.encode([input_text])[0]
-    response = ""
-
-    if index.ntotal > 0:
-        D, I = index.search(np.array([vec]), min(5, index.ntotal))
-        matches = [memory_text[idx] for idx, dist in zip(I[0], D[0]) if idx != -1 and dist < 0.8]
-        if matches:
-            response += "🧠 Related memories:\n- " + "\n- ".join(matches[:3]) + "\n\n"
-
-    # Store current memory
-    index.add(np.array([vec]))
-    memory_text.append(input_text)
-
-    # 🔥 Query LLM (Fireworks)
-    llm_response = query_fireworks(input_text)
-    response += f"🤖 Response:\n{llm_response}"
-    return response
-
-# === Gradio UI ===
-gradio_ui = gr.Interface(
-    fn=autonomous_agent,
-    inputs="text",
-    outputs="text",
-    title="Autonomous AI Agent",
-    description="Self-enhancing chatbot with memory + Fireworks LLM",
-)
-
-# === FastAPI App ===
-app = FastAPI()
-
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"], allow_methods=["*"], allow_headers=["*"],
-)
-
-# === Root Web UI with embedded Gradio ===
-@app.get("/", response_class=HTMLResponse)
-async def root():
-    return """
-    <html>
-        <head><title>Autonomous AI Agent</title></head>
-        <body>
-            <h2>Autonomous AI Agent</h2>
-            <iframe src="/gradio" width="100%" height="600"></iframe>
-        </body>
-    </html>
-    """
-
-# === Mount Gradio & static files ===
-app.mount("/gradio", gradio_ui.app)
-app.mount("/static", StaticFiles(directory="static"), name="static")
-
-# === For Hugging Face Spaces ===
-def get_app():
-    return app
+print("Sentence embeddings:")
+print(sentence_embeddings)