Update main.py

main.py

@@ -6,17 +6,18 @@ from transformers import pipeline, TextGenerationPipeline
 
 app = FastAPI()
 
-# Get Gemini API key from environment variable
-gemini_api_key = os.environ.get("GEMINI_API_KEY")
-if not gemini_api_key:
-    raise ValueError("GEMINI_API_KEY environment variable is not set.")
+# 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.")
 
 # Load the Gemini model using Hugging Face's pipeline
-# Make sure to use a model you have access to
+# Make sure to use a model you have access to and is available on Hugging Face
 generator: TextGenerationPipeline = pipeline(
     "text-generation",
-    model="google/gemma-2-2b-it",  # Replace if needed
-)  # IMPORTANT: **DO NOT** set `use_auth_token` here
+    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
+)  
 
 # Data model for the request body
 class Item(BaseModel):
@@ -31,14 +32,11 @@ async def generate_text(item: Item):
         if not item.prompt:
             raise HTTPException(status_code=400, detail="`prompt` field is required")
 
-        # Set API key in the headers BEFORE calling the pipeline
-        generator.model.config.use_auth_token = gemini_api_key  # Set the API key here
         output = generator(
             item.prompt,
             temperature=item.temperature,
             max_length=item.max_new_tokens, 
         )
-        generator.model.config.use_auth_token = None  # Reset after use
 
         return {"generated_text": output[0]['generated_text']}