dodany InferenceClient

app.py

@@ -1,69 +1,45 @@
-import os
 import spaces
+from huggingface_hub import InferenceClient
 import gradio as gr
-import transformers
-import torch
-from huggingface_hub import login
-import logging
-
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-model_id = "meta-llama/Meta-Llama-3.1-8B"
-
-def create_pipeline():
-    login(token=os.environ.get("MY_API_LLAMA_3_1"))
-    logger.info("Login successful")
-
-    config = transformers.AutoConfig.from_pretrained(model_id)
-    model = transformers.AutoModelForCausalLM.from_pretrained(
-        model_id,
-        config=config,
-        torch_dtype=torch.bfloat16,
-        device_map="auto",
-        use_cache=False
-    )
-    model.tie_weights()
-    tokenizer = transformers.AutoTokenizer.from_pretrained(model_id)
+import os
 
-    return transformers.pipeline(
-        "text-generation",
-        model=model,
-        tokenizer=tokenizer,
-        torch_dtype=torch.bfloat16,
-        device_map="auto"
-    )
+# Inicjalizacja klienta
+client = InferenceClient(
+    model='meta-llama/Meta-Llama-3.1-8B',
+    token=os.environ.get("MY_API_LLAMA_3_1")
+)
 
 @spaces.GPU(duration=60)
 def generate_response(chat, kwargs):
-    pipeline = create_pipeline()
-    try:
-        output = pipeline(chat, **kwargs)[0]['generated_text']
-        if output.endswith("</s>"):
-            output = output[:-4]
-        return output
-    except Exception as e:
-        logger.error(f"Error generating response: {str(e)}")
-        return f"Wystąpił błąd podczas generowania odpowiedzi: {str(e)}"
+    output = ''
+    stream = client.text_generation(chat, **kwargs, stream=True, details=True, return_full_text=False)
+    for response in stream:
+        output += response.token.text
+    if output.endswith("</s>"):  # Sprawdzamy, czy odpowiedź kończy się tagiem </s>
+        output = output[:-4]  # Usuwamy tag </s> z końca odpowiedzi
+    return output
 
 def function(prompt, history=[]):
     chat = "<s>"
     for user_prompt, bot_response in history:
         chat += f"[INST] {user_prompt} [/INST] {bot_response}</s> <s>"
-    chat += f"[INST] {prompt} [/INST]"
-
+    chat += f"[INST] {prompt} [/INST]"  # Zostawiamy tylko tag otwierający <s> na początku i kończymy ciąg zwykłym znacznikiem
     kwargs = dict(
-        max_new_tokens=4096,
-        do_sample=True,
         temperature=0.5,
+        max_new_tokens=4096,
         top_p=0.95,
         repetition_penalty=1.0,
+        do_sample=True,
         seed=1337
     )
 
-    return generate_response(chat, kwargs)
+    try:
+        output = generate_response(chat, kwargs)
+        return output
+    except Exception as e:
+        print(f"Error: {str(e)}")
+        return ''
 
-# Interfejs Gradio
 interface = gr.ChatInterface(
     fn=function,
     chatbot=gr.Chatbot(
@@ -85,11 +61,4 @@ interface = gr.ChatInterface(
     clear_btn=None
 )
 
-# API endpoint
-def api_predict(prompt):
-    return function(prompt)
-
-interface.launch(show_api=True, share=True)
-
-# Dodanie endpointu API
-gr.Interface(fn=api_predict, inputs="text", outputs="text").launch(share=True)
+interface.launch(show_api=True, share=True)