init

.idea/vcs.xml

@@ -0,0 +1,6 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="VcsDirectoryMappings">
+    <mapping directory="" vcs="Git" />
+  </component>
+</project>

app.py

@@ -1,29 +1,23 @@
-from flask import Flask, request, jsonify
+from fastapi import FastAPI
+from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
 
-app = Flask(__name__)
+app = FastAPI()
 
-MODEL_NAME = "meta-llama/Llama-2-7b-chat-hf"  # Ensure you have access!
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, use_auth_token=True)
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_NAME, torch_dtype=torch.float16, device_map="auto", use_auth_token=True
-)
+# Load Llama 2 model
+MODEL_NAME = "meta-llama/Llama-2-7b-chat-hf"
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype=torch.float16, device_map="auto")
 
-@app.route('/generate', methods=['POST'])
-def generate():
-    user_request = request.json.get("query")
-    if not user_request:
-        return jsonify({"error": "No query provided"}), 400
 
-    mongo_query = generate_mongo_query(user_request)
-    return jsonify({"mongo_query": mongo_query})
+@app.post("/generate")
+async def generate_text(data: dict):
+    prompt = data.get("prompt", "")
+    if not prompt:
+        return {"error": "No prompt provided"}
 
-def generate_mongo_query(user_request):
-    prompt = f"Convert this request to a MongoDB query: {user_request}"
-    inputs = tokenizer(prompt, return_tensors="pt").to("cuda" if torch.cuda.is_available() else "cpu")
-    output = model.generate(**inputs, max_length=150)
-    return tokenizer.decode(output[0], skip_special_tokens=True)
+    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
+    output = model.generate(**inputs, max_length=200)
+    response = tokenizer.decode(output[0], skip_special_tokens=True)
 
-if __name__ == '__main__':
-    app.run(host='0.0.0.0', port=7860)
+    return {"generated_text": response}

dockerfile

@@ -1,17 +1,23 @@
-# Use an official lightweight Python image
-FROM python:3.10
+# Use a Python base image
+FROM python:3.9
+
+# Create a user and set up environment
+RUN useradd -m -u 1000 user
+USER user
+ENV PATH="/home/user/.local/bin:$PATH"
 
 # Set working directory
 WORKDIR /app
 
-# Copy files to container
-COPY app.py requirements.txt /app/
+# Copy and install dependencies
+COPY --chown=user ./requirements.txt requirements.txt
+RUN pip install --no-cache-dir --upgrade -r requirements.txt
 
-# Install dependencies
-RUN pip install --no-cache-dir -r requirements.txt
+# Copy the application code
+COPY --chown=user . /app
 
-# Expose port for Flask
+# Expose port (optional, usually HF Spaces handles this)
 EXPOSE 7860
 
-# Run the Flask app
-CMD ["python", "app.py"]
+# Run the application
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]

requirements.txt

@@ -1,6 +1,5 @@
-flask
-torch
 transformers
+torch
 accelerate
-sentencepiece
-huggingface_hub
+fastapi
+uvicorn