Update main.py

main.py

@@ -1,10 +1,9 @@
 import os
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import pipeline
 from huggingface_hub import login
 
-
 # Define the FastAPI app
 app = FastAPI()
 
@@ -12,27 +11,18 @@ class InputText(BaseModel):
     input_text: str
     temperature: float = 1.0  # Default temperature
 
-HF_TOKEN = os.getenv("HF_TOKEN")
-    
+# 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)
 
-# Path to the pretrained model and tokenizer
-tokenizer = AutoTokenizer.from_pretrained("S1mp1eXXX/Mia_astral-1B")
-model = AutoModelForCausalLM.from_pretrained("S1mp1eXXX/Mia_astral-1B")
+# 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):
-    input_ids = tokenizer.encode(prompt, return_tensors='pt')
-    outputs = model.generate(
-        input_ids,
-        max_length=max_length,
-        repetition_penalty=1.2,
-        do_sample=True,
-        top_k=50,
-        top_p=0.95,
-        temperature=temperature,  # Set the temperature parameter
-        num_return_sequences=1
-    )
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+    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):