Update app.py

app.py

@@ -2,28 +2,28 @@ import os
 os.system("pip3 install transformers")
 os.system("pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu")
 os.system("pip3 install tensorflow")
-os.system("pip3 install bitsandbytes accelerate")
 
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, pipeline
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 import torch
 
-bnb_config = BitsAndBytesConfig(
-    load_in_4bit=True,
-    bnb_4bit_quant_type="nf4",
-    bnb_4bit_compute_dtype=torch.float16,
-    bnb_4bit_use_double_quant=True
-)
-
+def load_model():
+    model = AutoModelForCausalLM.from_pretrained(
+        "nvidia/Llama-3.1-Nemotron-Nano-8B-v1",
+        load_in_8bit=True,
+        device_map="auto", 
+        torch_dtype=torch.float16,
+        trust_remote_code=True
+    )
 
-model = AutoModelForCausalLM.from_pretrained(
-    "nvidia/Llama-3.1-Nemotron-Nano-8B-v1",
-    quantization_config=bnb_config,
-    device_map="auto",
-    trust_remote_code=True
-)
-tokenizer = AutoTokenizer.from_pretrained("nvidia/Llama-3.1-Nemotron-Nano-8B-v1")
+    tokenizer = AutoTokenizer.from_pretrained(
+        "nvidia/Llama-3.1-Nemotron-Nano-8B-v1",
+        trust_remote_code=True
+    )
+    
+    return model, tokenizer
 
+model, tokenizer = load_model()
 pipe = pipeline(
     "text-generation",
     model=model,
@@ -31,27 +31,42 @@ pipe = pipeline(
     device_map="auto"
 )
 
-def textgen(request):
-    messages = [
-        {"role": "user", "content": str(request)},
-    ]
-    
-    outputs = pipe(
-        messages,
-        max_new_tokens=512,
-        do_sample=True,
-        temperature=0.7,
-        top_p=0.9
-    )
+def generate_response(request):
+    try:
+        messages = [
+            {"role": "user", "content": str(request)},
+        ]
+        
+        outputs = pipe(
+            messages,
+            max_new_tokens=512,
+            do_sample=True,
+            temperature=0.7,
+            top_p=0.9,
+            repetition_penalty=1.1
+        )
+        
+        return outputs[0]["generated_text"][-1]['content']
     
-    return outputs[0]["generated_text"][-1]['content']
+    except Exception as e:
+        return f"Произошла ошибка: {str(e)}"
 
 demo = gr.Interface(
-    fn=textgen,
-    inputs=gr.Textbox(label="Ваш запрос", placeholder="Введите ваш вопрос здесь..."),
-    outputs=gr.Textbox(label="Ответ модели"),
-    title="Chat with Llama-3.1-Nemotron-Nano (4-bit quantized)",
-    description="Квантованная 4-bit версия модели NVIDIA Llama-3.1-Nemotron-Nano-8B"
+    fn=generate_response,
+    inputs=gr.Textbox(
+        label="Ваш запрос",
+        placeholder="Введите ваш вопрос здесь...",
+        lines=3
+    ),
+    outputs=gr.Textbox(
+        label="Ответ модели",
+        lines=5
+    ),
+    title="Chat with 8-bit Llama-3.1-Nemotron-Nano",
+    description="8-битная квантованная версия модели NVIDIA Llama-3.1-Nemotron-Nano-8B",
+    allow_flagging="never"
 )
 
-demo.launch(share=True)
+# Запускаем интерфейс
+if __name__ == "__main__":
+    demo.launch(share=True)