Update app.py

app.py

@@ -1,9 +1,7 @@
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
-from huggingface_hub import login
+import requests
 import os
-import torch
 import logging
 
 # Read the Hugging Face token from the environment variable
@@ -11,9 +9,6 @@ token = os.getenv("HUGGING_FACE_HUB_TOKEN")
 if token is None:
     raise ValueError("Hugging Face token not found in environment variables. Please set the HUGGING_FACE_HUB_TOKEN secret in Hugging Face Spaces.")
 
-# Log in with the token
-login(token=token)
-
 # Initialize FastAPI app
 app = FastAPI()
 
@@ -21,28 +16,9 @@ app = FastAPI()
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# Load a smaller model (GPT-Neo-125M) and tokenizer
-model_id = "EleutherAI/gpt-neo-125M"  # Use GPT-Neo-125M for faster performance
-tokenizer = AutoTokenizer.from_pretrained(model_id)
-
-# Set pad_token if it doesn't exist
-if tokenizer.pad_token is None:
-    tokenizer.pad_token = tokenizer.eos_token  # Use eos_token as pad_token
-
-# Load the model without quantization for CPU
-logger.info("Loading model...")
-model = AutoModelForCausalLM.from_pretrained(
-    model_id,
-    torch_dtype=torch.float32,  # Use FP32 for CPU compatibility
-    device_map="auto"           # Automatically offload to available devices
-)
-
-# Create a text generation pipeline
-pipe = pipeline(
-    "text-generation",
-    model=model,
-    tokenizer=tokenizer
-)
+# Hugging Face Inference API endpoint for BLOOM-7B
+API_URL = "https://api-inference.huggingface.co/models/bigscience/bloom-7b1"  # Use BLOOM-7B
+headers = {"Authorization": f"Bearer {token}"}
 
 # Define request body schema
 class TextGenerationRequest(BaseModel):
@@ -59,17 +35,25 @@ async def generate_text(request: TextGenerationRequest):
     try:
         logger.info("Generating text...")
         
-        # Generate text using the pipeline with the user's prompt
-        outputs = pipe(
-            request.prompt,  # Use the user's prompt directly
-            max_new_tokens=request.max_new_tokens,
-            temperature=request.temperature,
-            top_k=request.top_k,
-            top_p=request.top_p,
-            do_sample=request.do_sample,
-            return_full_text=False  # Exclude the input prompt from the output
-        )
-        return {"generated_text": outputs[0]["generated_text"]}
+        # Prepare the payload for the Hugging Face Inference API
+        payload = {
+            "inputs": request.prompt,
+            "parameters": {
+                "max_new_tokens": request.max_new_tokens,
+                "temperature": request.temperature,
+                "top_k": request.top_k,
+                "top_p": request.top_p,
+                "do_sample": request.do_sample,
+            },
+        }
+
+        # Send request to the Hugging Face Inference API
+        response = requests.post(API_URL, headers=headers, json=payload)
+        response.raise_for_status()  # Raise an error for bad responses (4xx or 5xx)
+
+        # Extract the generated text from the response
+        generated_text = response.json()[0]["generated_text"]
+        return {"generated_text": generated_text}
     except Exception as e:
         logger.error(f"Error generating text: {e}")
         raise HTTPException(status_code=500, detail=str(e))