Update main.py

main.py

@@ -5,63 +5,67 @@ from transformers import pipeline
 from pydantic import BaseModel
 from typing import Optional, Any
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, TextStreamer, GenerationConfig
 
-app = FastAPI()
+# Authentication
+from huggingface_hub import interpreter_login
+interpreter_login()
+
+# Packages and model loading
+import torch
+base_model_id = "152334H/miqu-1-70b-sf"
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_compute_dtype=torch.bfloat16
+)
 
-# Initialize device
-def initialize_device():
-    global device
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+base_model = AutoModelForCausalLM.from_pretrained(
+    base_model_id,
+    quantization_config=bnb_config,
+    device_map="auto",
+    trust_remote_code=True,
+)
 
-initialize_device()
+# Tokenizer loading
+eval_tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-70b-hf", add_bos_token=True, trust_remote_code=True, use_auth_token=True)
+
+# Streamer
+streamer = TextStreamer(eval_tokenizer)
+
+# App definition
+app = FastAPI()
 
 # Helper function to read raw request bodies
 async def parse_raw(request: Request):
     return await request.body()
 
-# Initialize the model and tokenizer with the corrected pre-trained weights
-def init_corrected_model():
-    global model_config, model, tokenizer
+# Generate text
+def generate_text(prompt: str) -> str:
+    model_input = eval_tokenizer(prompt, return_tensors="pt").to("cuda")
 
-    try:
-        model_config = AutoConfig.from_pretrained("152334H/miqu-1-70b-sf", trust_remote_code=True)
-        model = AutoModelForCausalLM.from_pretrained("152334H/miqu-1-70b-sf", config=model_config).to(device)
-        tokenizer = AutoTokenizer.from_pretrained("152334H/miqu-1-70b-sf")
-    except Exception as e:
-        print("[WARNING]: Failed to load model and tokenizer conventionally.")
-        print(f"Exception: {e}")
+    base_model.eval()
+    with torch.no_grad():
+        generated_sequences = base_model.generate(
+            **model_input,
+            max_new_tokens=4096,
+            repetition_penalty=1.1,
+            do_sample=True,
+            temperature=1,
+            streamer=streamer,
+        )
 
-        model_config = AutoConfig.from_pretrained("152334H/miqu-1-70b-sf", trust_remote_code=True)
-        model = AutoModelForCausalLM.from_pretrained("152334H/miqu-1-70b-sf", config=model_config).to(device)
-        tokenizer = AutoTokenizer.from_pretrained("152334H/miqu-1-70b-sf")
+    return eval_tokenizer.decode(generated_sequences[0], skip_special_tokens=True)
 
-init_corrected_model()
-
-# Utility function to generate answers from the model
-def miuk_answer(query: str) -> dict[str, 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 {"output": tokenizer.decode(answer[:, 0])}
-
-# Endpoint handler to receive incoming queries and pass them to the utility function for processing
-@app.post("/infer_miku")
-async def infer_endpoint(data: BaseModel = Depends(parse_raw)):
+# Route for generating text
+@app.post("/generate_text")
+async def generate_text_route(data: BaseModel = Depends(parse_raw)):
     input_text = data.raw.decode("utf-8")
-
-    if input_text is None or len(input_text) == 0:
+    if not input_text or len(input_text) <= 0:
         return JSONResponse({"error": "Empty input received."}, status_code=400)
 
-    result = miuk_answer(input_text)
-    return result
-
-@app.get("/infer_miku")
-def get_default_inference_endpoint():
-    return {"message": "Use POST method to submit input data"}
+    return {"output": generate_text(input_text)}
 
 # Mount static files
-app.mount("/static", StaticFiles(directory="static"), name="static")
-
-# Initialization done
-print("Initialization Complete.")
+app.mount("/static", StaticFiles(directory="static"), name="static")