Update Dockerfile and FastAPI app

Dockerfile

@@ -1,21 +1,23 @@
+# Use the official Hugging Face TGI Docker image as the base
 FROM huggingface/transformers-pytorch-gpu:latest
 
+
 RUN pip install --upgrade pip
 RUN pip install transformers torch fastapi uvicorn
 
 
-ENV TRANSFORMERS_CACHE=/tmp/huggingface_cache
-
-
-ENV MODEL_NAME="your-username/your-finetuned-model"
+ENV MODEL_NAME="Hadeel11/fine-tuned-model"
 ENV USE_FP16=True
 
+
 COPY app.py /app/app.py
 
 WORKDIR /app
 
+
 EXPOSE 8080
 
+
 CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "8080"]
 
 

app.py

@@ -4,9 +4,9 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 
 app = FastAPI()
 
-model_name = "Hadeel11/fine-tuned-model"
-model = AutoModelForCausalLM.from_pretrained(model_name)
-tokenizer = AutoTokenizer.from_pretrained(model_name)
+# Load your fine-tuned model and tokenizer
+model = AutoModelForCausalLM.from_pretrained("Hadeel11/fine-tuned-model")
+tokenizer = AutoTokenizer.from_pretrained("Hadeel11/fine-tuned-model")
 
 class Query(BaseModel):
     question: str
@@ -18,3 +18,4 @@ async def predict(query: Query):
     outputs = model.generate(**inputs)
     answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
     return {"answer": answer}
+