Upload 2 files

Dockerfile

@@ -1,13 +1,21 @@
-FROM python:3.10-slim
+FROM python:3.10
 
-RUN apt-get update && apt-get install -y git && rm -rf /var/lib/apt/lists/*
+# Set working directory
+WORKDIR /app
 
-RUN pip install --upgrade pip
-RUN pip install torch transformers fastapi uvicorn bitsandbytes accelerate
+# Copy local files
+COPY . .
 
-WORKDIR /app
-COPY main.py .
-ENV HF_HOME=/app/hf_cache
-RUN mkdir -p /app/hf_cache
+# Install Python dependencies
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Set Hugging Face cache directory to writable one
+ENV HF_HOME=/data
+RUN mkdir -p /data && chmod 777 /data
 
+# Expose the port
+EXPOSE 7860
+
+# Run the FastAPI app with Uvicorn
 CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "7860"]
+

main.py

@@ -1,25 +1,27 @@
-from fastapi import FastAPI, Request
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
-import torch
 import os
+from fastapi import FastAPI
+from pydantic import BaseModel
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 
-os.environ["HF_HOME"] = "/app/hf_cache"
-
-app = FastAPI()
+# Force Hugging Face cache to a writable dir
+os.environ["HF_HOME"] = "/data"
 
 model_name = "howtomakepplragequit/phi2-lora-instruct"
+
+# Load tokenizer and model
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(
-    model_name,
-    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-    device_map="auto"
-)
-pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
+model = AutoModelForCausalLM.from_pretrained(model_name)
+
+# Create pipeline
+generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
+
+# FastAPI app setup
+app = FastAPI()
+
+class Prompt(BaseModel):
+    prompt: str
 
 @app.post("/generate")
-async def generate(request: Request):
-    data = await request.json()
-    prompt = data.get("prompt", "")
-    formatted = f"### Instruction:\n{prompt}\n\n### Response:\n"
-    result = pipe(formatted, max_new_tokens=200)[0]["generated_text"]
-    return {"response": result.split("### Response:")[-1].strip()}
+def generate_text(data: Prompt):
+    output = generator(data.prompt, max_length=200, do_sample=True)[0]["generated_text"]
+    return {"response": output}