Update main.py

main.py

@@ -4,9 +4,23 @@ from fastapi.staticfiles import StaticFiles
 from transformers import pipeline
 from pydantic import BaseModel
 from typing import Union
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
 app = FastAPI()
 
+# Load the MIKU model and tokenizer
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = AutoModelForCausalLM.from_pretrained("miqudev/miqu-1-70b").to(device)
+tokenizer = AutoTokenizer.from_pretrained("miqudev/miqu-1-70b")
+
+def miuk_answer(query: str):
+    query_tokens = tokenizer.encode(query, return_tensors="pt")
+    query_tokens = query_tokens.to(device)
+    answer = model.generate(query_tokens, max_length=128, temperature=1, pad_token_id=tokenizer.pad_token_id)
+    return tokenizer.decode(answer[:, 0]).replace(" ", "")
+
+
 # Serve the static files
 app.mount("/static", StaticFiles(directory="static"), name="static")
 
@@ -41,16 +55,6 @@ def get_default_inference_endpoint():
 def index():
     return './static/index.html'
 
-from typing import Union
-from transformers import pipeline
-
-# Load miku pipeline
-pipe_miku = pipeline("text-generation", model="miqudev/miqu-1-70b")
-
-def miku(input):
-    output = pipe_miku(input)
-    return {"output": output[0]["generated_text"]}
-
 @app.post("/infer_miku")
 def infer_endpoint(data: dict):
     """Receive input and generate text."""
@@ -62,10 +66,12 @@ def infer_endpoint(data: dict):
     if input_text is None:
         return {"error": "No input text detected."}
     else:
-        result = miku(input_text)
+        result = {"output": miuk_answer(input_text)}
         return result
 
 @app.get("/infer_miku")
 def get_default_inference_endpoint():
     return {"message": "Use POST method to submit input data"}
+    
+