Install transformers dependency (#1)

- Install transformers dependency (f274ca866a4c5cdfd36394ed24e858445ab4893d)
- Update python script (231f15f5049f32e1080c2b1aeb6a2c8182964302)

app.py

@@ -1,10 +1,31 @@
 import gradio as gr
-from huggingface_hub import InferenceClient
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
-"""
-For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
-"""
-client = InferenceClient("AIDC-AI/Marco-o1")
+tokenizer = AutoTokenizer.from_pretrained("AIDC-AI/Marco-o1")
+model = AutoModelForCausalLM.from_pretrained("AIDC-AI/Marco-o1")
+
+def load_model_and_tokenizer(path):
+    tokenizer = AutoTokenizer.from_pretrained(path, trust_remote_code=True)
+    model = AutoModelForCausalLM.from_pretrained(path, trust_remote_code=True).to('cuda:0')
+    model.eval()
+    return tokenizer, model
+
+def generate_response(model, tokenizer,
+                      input_ids, attention_mask,
+                      max_new_tokens=4096):
+    generated_ids = input_ids
+    with torch.inference_mode():
+        for _ in range(max_new_tokens):
+            outputs = model(input_ids=generated_ids, attention_mask=attention_mask)
+            next_token_id = torch.argmax(outputs.logits[:, -1, :], dim=-1).unsqueeze(-1)
+            generated_ids = torch.cat([generated_ids, next_token_id], dim=-1)
+            attention_mask = torch.cat([attention_mask, torch.ones_like(next_token_id)], dim=-1)
+            new_token = tokenizer.decode(next_token_id.squeeze(), skip_special_tokens=True)
+            print(new_token, end='', flush=True)
+            if next_token_id.item() == tokenizer.eos_token_id:
+                break
+    return tokenizer.decode(generated_ids[0][input_ids.shape[-1]:], skip_special_tokens=True)
 
 
 def respond(
@@ -25,20 +46,9 @@ def respond(
 
     messages.append({"role": "user", "content": message})
 
-    response = ""
-
-    for message in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p,
-    ):
-        token = message.choices[0].delta.content
-
-        response += token
-        yield response
-
+    text = tokenizer.apply_chat_template(history, tokenize=False, add_generation_prompt=True)
+    model_inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=4096).to('cuda:0')
+    yield generate_response(model, tokenizer, model_inputs.input_ids, model_inputs.attention_mask)
 
 """
 For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface

requirements.txt

@@ -1 +1 @@
-huggingface_hub==0.25.2
+transformers==4.46.3