poprawki

app.py

@@ -2,29 +2,49 @@ import os
 import spaces
 import gradio as gr
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
+from threading import Thread
+from queue import Queue
 
 model_id = "meta-llama/Meta-Llama-3.1-8B"
+tokenizer = AutoTokenizer.from_pretrained(model_id, token=os.environ.get("MY_API_LLAMA_3_1"))
 
-@spaces.GPU(duration=60)
-def load_model():
-    tokenizer = AutoTokenizer.from_pretrained(model_id, token=os.environ.get("MY_API_LLAMA_3_1"))
-    model = AutoModelForCausalLM.from_pretrained(
-        model_id,
-        token=os.environ.get("MY_API_LLAMA_3_1"),
-        torch_dtype=torch.bfloat16,
-        device_map="auto",
-        low_cpu_mem_usage=True
-    )
-    return pipeline("text-generation", model=model, tokenizer=tokenizer)
+model = None
+model_load_queue = Queue()
 
-pipe = load_model()
+def load_model():
+    global model
+    if model is None:
+        model = AutoModelForCausalLM.from_pretrained(
+            model_id,
+            token=os.environ.get("MY_API_LLAMA_3_1"),
+            torch_dtype=torch.bfloat16,
+            device_map="auto",
+            low_cpu_mem_usage=True
+        )
+    model_load_queue.put(model)
 
 @spaces.GPU(duration=60)
 def generate_response(chat, kwargs):
-    output = pipe(chat, **kwargs)[0]['generated_text']
-    if output.endswith("</s>"):
-        output = output[:-4]
+    global model
+    if model is None:
+        Thread(target=load_model).start()
+        model = model_load_queue.get()
+
+    inputs = tokenizer(chat, return_tensors="pt").to(model.device)
+    streamer = TextIteratorStreamer(tokenizer, timeout=10., skip_prompt=True, skip_special_tokens=True)
+
+    generation_kwargs = dict(inputs, streamer=streamer, **kwargs)
+    thread = Thread(target=model.generate, kwargs=generation_kwargs)
+    thread.start()
+
+    output = ""
+    for new_text in streamer:
+        output += new_text
+        if output.endswith("</s>"):
+            output = output[:-4]
+            break
+
     return output
 
 def function(prompt, history=[]):
@@ -33,11 +53,11 @@ def function(prompt, history=[]):
         chat += f"[INST] {user_prompt} [/INST] {bot_response}</s> <s>"
     chat += f"[INST] {prompt} [/INST]"
     kwargs = dict(
-        temperature=0.5,
         max_new_tokens=4096,
+        do_sample=True,
+        temperature=0.5,
         top_p=0.95,
         repetition_penalty=1.0,
-        do_sample=True,
         seed=1337
     )