Update main.py

main.py

@@ -6,17 +6,16 @@ from transformers import pipeline, TextGenerationPipeline
 
 app = FastAPI()
 
-# Get Gemini API key from environment variable (NOT USED FOR HUGGING FACE AUTH)
-# gemini_api_key = os.environ.get("GEMINI_API_KEY")
-# if not gemini_api_key:
-#     raise ValueError("GEMINI_API_KEY environment variable is not set.")
+# Ensure your Hugging Face token is set as an environment variable
+huggingface_token = os.environ.get("TOKEN")
+if not huggingface_token:
+    raise ValueError("TOKEN environment variable is not set.")
 
-# Load the Gemini model using Hugging Face's pipeline
-# Make sure to use a model you have access to and is available on Hugging Face
+# Load the model using Hugging Face's pipeline with the token
 generator: TextGenerationPipeline = pipeline(
     "text-generation",
     model="google/gemma-2-2b-it",  # The model you want to use
-    use_auth_token= os.environ.get("TOKEN")  # The Hugging Face token you got after login
+    use_auth_token=huggingface_token  # Use the Hugging Face token for authentication
 )  
 
 # Data model for the request body
@@ -44,4 +43,4 @@ async def generate_text(item: Item):
         raise HTTPException(status_code=500, detail=f"An error occurred: {e}")
 
 if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=8000)
+    uvicorn.run(app, host="0.0.0.0", port=8000)