main.py

import os
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from transformers import pipeline
from huggingface_hub import login

# Define the FastAPI app
app = FastAPI()

class InputText(BaseModel):
    input_text: str
    temperature: float = 1.0  # Default temperature

# Hugging Face authentication
HF_TOKEN = os.getenv("HF_TOKEN")  # Replace with your actual token or set it as an environment variable

# Login to Hugging Face
login(token=HF_TOKEN)

# Define the text generation pipeline
text_generation_pipeline = pipeline("text-generation", model="S1mp1eXXX/Mia_astral-1B", use_auth_token=HF_TOKEN)

def generate_text(prompt, max_length=1400, temperature=1.0):
    outputs = text_generation_pipeline(prompt, max_length=max_length, temperature=temperature, num_return_sequences=1)
    return outputs[0]['generated_text']

@app.post("/generate-text/")
async def generate_text_post(data: InputText):
    generated_text = generate_text(data.input_text, max_length=140, temperature=data.temperature)
    return {"generated_text": generated_text}

@app.get("/generate-text/")
async def generate_text_get(input_text: str, temperature: float = 1.0):
    generated_text = generate_text(input_text, max_length=480, temperature=temperature)
    return {"generated_text": generated_text}

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=5000)