Update main.py

main.py

@@ -15,9 +15,8 @@ if not gemini_api_key:
 # Make sure to use a model you have access to
 generator: TextGenerationPipeline = pipeline(
     "text-generation",
-    model="llama-duo/gemma7b-summarize-gemini1.5flash-80k",  # Replace if needed
-    use_auth_token=gemini_api_key
-)
+    model="gemini-1.5-flash",  # Replace if needed
+)  # IMPORTANT: **DO NOT** set `use_auth_token` here
 
 # Data model for the request body
 class Item(BaseModel):
@@ -32,11 +31,14 @@ 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']}