Upload 3 files

Dockerfile

@@ -1,20 +1,20 @@
-# Use the official PyTorch image with CUDA support
-FROM pytorch/pytorch:2.0.0-cuda11.7-cudnn8-runtime
+# Use an official Python runtime as a parent image
+FROM python:3.9-slim
 
-# Install git to clone the model repository
-RUN apt-get update && apt-get install -y git
+# Set the working directory
+WORKDIR /app
 
-# Install other dependencies
-RUN pip install transformers
+# Copy the current directory contents into the container at /app
+COPY . /app
 
-# Clone the unsloth/llama-3-8b-bnb-4bit model repository
-RUN git clone https://huggingface.co/unsloth/llama-3-8b-bnb-4bit
+# Install any needed packages specified in requirements.txt
+RUN pip install --no-cache-dir -r requirements.txt
 
-# Set the working directory to the model directory
-WORKDIR /unsloth/llama-3-8b-bnb-4bit
+# Make port 80 available to the world outside this container
+EXPOSE 80
 
-# Copy the inference script
-COPY inference.py .
+# Define environment variable
+ENV NAME World
 
-# Run the inference script
-CMD ["python", "inference.py"]
+# Run app.py when the container launches
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "80"]

app.py

@@ -0,0 +1,21 @@
+from fastapi import FastAPI
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
+
+app = FastAPI()
+
+# Load the model and tokenizer
+model_name = "unsloth/llama-3-8b-bnb-4bit"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name)
+
+@app.get("/")
+def read_root():
+    return {"Hello": "World"}
+
+@app.post("/generate/")
+def generate(prompt: str):
+    inputs = tokenizer(prompt, return_tensors="pt")
+    outputs = model.generate(inputs["input_ids"], max_length=50)
+    text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return {"generated_text": text}

requirements.txt

@@ -0,0 +1,4 @@
+fastapi
+uvicorn
+torch
+transformers