Add Space config for Hugging Face5

README.md

@@ -1,9 +1,15 @@
 ---
-title: Capybara Classifier 🦫
+title: Capybara Classifier
 emoji: 🦫
-colorFrom: green
+colorFrom: gray
 colorTo: blue
-
+sdk: gradio
 app_file: app.py
 pinned: false
 ---
+
+Capybara-3B модель для классификации текста как `запрос` или `реклама`.
+
+Пример:  
+`Сколько стоит консультация?` → **запрос**  
+`Скидка 50% на айфон!` → **реклама**

app.py

@@ -1,37 +1,48 @@
-# app.py
 import os
+import gradio as gr
+import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
-from fastapi import FastAPI
-from pydantic import BaseModel
 
-hf_token = os.getenv("HF_TOKEN")
-model_path = "./capybara-finetuned"  # локальная модель или путь на Hugging Face
+# Токен, если модель приватная
+hf_token = os.getenv("HF_TOKEN", None)
+model_path = "./capybara-finetuned"  # или HF-репозиторий, например: "NousResearch/Nous-Capybara-3B-V1.9"
 
+# Загружаем модель и токенизатор
 tokenizer = AutoTokenizer.from_pretrained(
     model_path,
     token=hf_token,
     trust_remote_code=True,
-    use_fast=False,
+    use_fast=True,
 )
 
 model = AutoModelForCausalLM.from_pretrained(
     model_path,
     token=hf_token,
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
     device_map="auto",
-    torch_dtype="auto",
-    trust_remote_code=True
+    trust_remote_code=True,
 )
 
+# Создаём пайплайн
 pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
 
-app = FastAPI()
+# Функция классификации
+def classify(text):
+    prompt = f"### Вопрос:\n{text}\n\n### Класс:"
+    try:
+        result = pipe(prompt, max_new_tokens=10, do_sample=False)[0]["generated_text"]
+        label = result.split("### Класс:")[-1].strip().split()[0].lower()
+        return f"🔍 Класс: **{label}**"
+    except Exception as e:
+        return f"❌ Ошибка: {str(e)}"
 
-class Input(BaseModel):
-    text: str
+# Интерфейс Gradio
+iface = gr.Interface(
+    fn=classify,
+    inputs=gr.Textbox(lines=4, placeholder="Введите сообщение..."),
+    outputs="markdown",
+    title="Capybara Text Classifier 🦫",
+    description="Классификация текста как 'запрос' или 'реклама' с помощью Capybara-3B",
+)
 
-@app.post("/classify")
-async def classify(input: Input):
-    prompt = f"### Вопрос:\n{input.text}\n\n### Класс:"
-    output = pipe(prompt, max_new_tokens=10, do_sample=False)[0]["generated_text"]
-    label = output.split("### Класс:")[-1].strip().split()[0].lower()
-    return {"label": label}
+iface.launch()

requirements.txt

@@ -1,7 +1,6 @@
-transformers==4.37.2
-accelerate==0.27.2
-torch>=2.1.0
-peft==0.8.2
-bitsandbytes==0.43.1
-uvicorn
-fastapi
+gradio
+torch
+transformers
+peft
+accelerate
+bitsandbytes