Dubeg

main.py

@@ -1,6 +1,12 @@
 from fastapi import FastAPI
 from typing import List
 from vllm import LLM, SamplingParams
+import os
+from dotenv import load_dotenv
+
+load_dotenv()
+
+token: str = os.environ.get("HUGGINGFACE_TOKEN")
 
 app = FastAPI()
 
@@ -8,10 +14,10 @@ app = FastAPI()
 @app.get("/llm_inference")
 def read_root(
         prompt: str,
-        model: str = "meta-llama/Llama-2-7B-hf",
+        model: str = "meta-llama/Llama-2-7b-hf",
         temperature: float = 0.,
         max_tokens: int = 1024) -> List:
-    sampling_params = SamplingParams(temperature=temperature, max_tokens=max_tokens)
+    sampling_params = SamplingParams(temperature=temperature, max_tokens=max_tokens, token=token)
 
     llm = LLM(model=model)