Update main.py

main.py

@@ -2,141 +2,40 @@ import os
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 import uvicorn
-import requests
-from google.auth.transport.requests import Request
-from google.oauth2 import id_token
+from transformers import pipeline, TextGenerationPipeline
 
 app = FastAPI()
 
-# Define the Gemini API endpoint for the primary model
-primary_url = "https://us-east4-aiplatform.googleapis.com/v1/projects/gen-lang-client-0770444709/locations/us-east4/publishers/google/models/gemini-1.5-flash-latest:generateContent"
+# Load the Gemini model using Hugging Face's pipeline
+generator: TextGenerationPipeline = pipeline(
+    "text-generation", 
+    model="google/gemini-1.5-flash"  # Or the specific Gemini model you have access to
+)
 
-# Define the data model for the request body
+# Define the data model for the request body (simplified for clarity)
 class Item(BaseModel):
-    input: str = None
-    system_prompt: str = None
-    system_output: str = None
-    history: list = None
-    templates: list = None
-    temperature: float = 0.0
-    max_new_tokens: int = 1048
-    top_p: float = 0.15
-    repetition_penalty: float = 1.0
-    key: str = None
-
-# Function to obtain an OAuth 2.0 access token
-def get_access_token():
-    auth_req = Request()
-    target_audience = primary_url.split("models")[0]  # Extract the base URL
-    token = id_token.fetch_id_token(auth_req, target_audience)
-    return token
-
-# Function to generate the response JSON
-def generate_response_json(item, output, tokens, model_name):
-    return {
-        "settings": {
-            "input": item.input if item.input is not None else "",
-            "system prompt": item.system_prompt if item.system_prompt is not None else "",
-            "system output": item.system_output if item.system_output is not None else "",
-            "temperature": f"{item.temperature}" if item.temperature is not None else "",
-            "max new tokens": f"{item.max_new_tokens}" if item.max_new_tokens is not None else "",
-            "top p": f"{item.top_p}" if item.top_p is not None else "",
-            "repetition penalty": f"{item.repetition_penalty}" if item.repetition_penalty is not None else "",
-            "do sample": "True",
-            "seed": "42"
-        },
-        "response": {
-            "output": output.strip().lstrip('\n').rstrip('\n').lstrip('').rstrip('').strip(),
-            "unstripped": output,
-            "tokens": tokens,
-            "model": "primary",
-            "name": model_name
-        }
-    }
-
-# Function to call the Gemini API
-def call_gemini_api(url, input_text, generate_kwargs):
-    access_token = get_access_token()
-    headers = {
-        "Authorization": f"Bearer {access_token}",
-        "Content-Type": "application/json"
-    }
-    data = {
-        "contents": [{"role": "user", "parts": [{"text": input_text}]}],
-        "generationConfig": {
-            "temperature": generate_kwargs['temperature'],
-            "maxOutputTokens": generate_kwargs['max_new_tokens'],
-            "topP": generate_kwargs['top_p'],
-            "repetitionPenalty": generate_kwargs['repetition_penalty'],
-        }
-    }
-    try:
-        response = requests.post(url, headers=headers, json=data)
-        response.raise_for_status()
-        return response.json()
-    except requests.exceptions.HTTPError as http_err:
-        raise HTTPException(status_code=response.status_code, detail=f"HTTP error occurred: {http_err}")
-    except Exception as err:
-        raise HTTPException(status_code=500, detail=f"An error occurred: {err}")
+    prompt: str
+    temperature: float = 0.7  # Adjust default as needed
+    max_new_tokens: int = 128  # Adjust default as needed
 
 # Endpoint for generating text
 @app.post("/")
-async def generate_text(item: Item = None):
+async def generate_text(item: Item):
     try:
-        if item is None:
-            raise HTTPException(status_code=400, detail="JSON body is required.")
-        if item.input is None and item.system_prompt is None or item.input == "" and item.system_prompt == "":
-            raise HTTPException(status_code=400, detail="Parameter `input` or `system prompt` is required.")
-
-        input_ = ""
-        if item.system_prompt is not None and item.system_output is not None:
-            input_ = f"<s>[INST] {item.system_prompt} [/INST] {item.system_output}</s>"
-        elif item.system_prompt is not None:
-            input_ = f"<s>[INST] {item.system_prompt} [/INST]</s>"
-        elif item.system_output is not None:
-            input_ = f"<s>{item.system_output}</s>"
-
-        if item.templates is not None:
-            for num, template in enumerate(item.templates, start=1):
-                input_ += f"\n<s>[INST] Beginning of archived conversation {num} [/INST]</s>"
-                for i in range(0, len(template), 2):
-                    input_ += f"\n<s>[INST] {template[i]} [/INST]"
-                    input_ += f"\n{template[i + 1]}</s>"
-                input_ += f"\n<s>[INST] End of archived conversation {num} [/INST]</s>"
-
-        input_ += f"\n<s>[INST] Beginning of active conversation [/INST]</s>"
-        if item.history is not None:
-            for input_text, output_text in item.history:
-                input_ += f"\n<s>[INST] {input_text} [/INST]"
-                input_ += f"\n{output_text}"
-        input_ += f"\n<s>[INST] {item.input} [/INST]"
-
-        temperature = float(item.temperature)
-        if temperature < 1e-2:
-            temperature = 1e-2
-        top_p = float(item.top_p)
-
-        generate_kwargs = dict(
-            temperature=temperature,
-            max_new_tokens=item.max_new_tokens,
-            top_p=top_p,
-            repetition_penalty=item.repetition_penalty,
-            do_sample=True,
-            seed=42,
+        if not item.prompt:
+            raise HTTPException(status_code=400, detail="`prompt` field is required")
+
+        # Call the Gemini model through Hugging Face's pipeline
+        output = generator(
+            item.prompt,
+            temperature=item.temperature,
+            max_length=item.max_new_tokens, 
         )
 
-        tokens = 0
-        response_data = call_gemini_api(primary_url, input_, generate_kwargs)
-        output = response_data['candidates'][0]['content']['parts'][0]['text']
-        tokens = response_data['usageMetadata']['totalTokenCount']
-
-        return generate_response_json(item, output, tokens, primary_url)
-
-    except HTTPException as http_error:
-        raise http_error
+        return {"generated_text": output[0]['generated_text']}
 
     except Exception as e:
         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)