RuBERT

app.py

@@ -1,39 +1,32 @@
 import gradio as gr
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig, BitsAndBytesConfig
 
-model_name = "ai-sage/GigaChat-20B-A3B-instruct-v1.5-int8"
+model_name = "DeepPavlov/rubert-base-cased"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForSequenceClassification.from_pretrained(model_name)
 
-bnb_config = BitsAndBytesConfig(
-    load_in_4bit=True,
-    bnb_4bit_quant_type="nf4",
-    bnb_4bit_compute_dtype=torch.bfloat16,
-    bnb_4bit_use_double_quant=True,
-    llm_int8_enable_fp32_cpu_offload=True
-)
-
-tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+texts = [
+    "Фильм был просто потрясающий, мне понравилось!",
+    "Это худший сервис, который я видел.",
+    "Нейтральное сообщение без эмоций."
+]
 
-device_map = {
-    "transformer.wte": "cpu",
-    "transformer.wpe": "cpu",
-    "transformer.h": "cpu",
-    "transformer.ln_f": "cpu",
-    "lm_head": "cpu",
-}
+inputs = tokenizer(texts, padding=True, truncation=True, return_tensors="pt", max_length=512)
 
-model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True, quantization_config=bnb_config, device_map="cpu")
-model.generation_config = GenerationConfig.from_pretrained(model_name)
+with torch.no_grad():
+    outputs = model(**inputs)
+    predictions = torch.softmax(outputs.logits, dim=1)
 
-messages = [
-    {"role": "user", "content": "Докажи теорему о неподвижной точке"}
-]
-input_tensor = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt")
-outputs = model.generate(input_tensor.to(model.device))
+labels = ["negative", "neutral", "positive"]
 
-result = tokenizer.decode(outputs[0][input_tensor.shape[1]:], skip_special_tokens=False)
+for text, pred in zip(texts, predictions):
+    print(f"Текст: {text}")
+    for i, score in enumerate(pred):
+        print(f"{labels[i]}: {score:.4f}")
+    print("---")
 
 with gr.Blocks() as demo:
-    gr.Textbox(result, label="Вывод строки", interactive=False)
+    gr.Textbox("СМОТРИ ЕПТА ЛОГИ", label="Вывод строки", interactive=False)
 
 demo.launch()