Locon213 commited on
Commit
4d3834d
·
verified ·
1 Parent(s): 9c4abee

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +8 -9
app.py CHANGED
@@ -1,20 +1,19 @@
1
- import gradio as gr
2
- from transformers import AutoModelForCausalLM, AutoTokenizer
3
- import torch
4
 
5
  # Имя модели
6
  model_name = "nvidia/Hymba-1.5B-Instruct"
7
 
8
- # Загрузка токенизатора и модели
9
- tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
10
- model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True)
11
 
12
  # Функция для генерации текста
13
  def generate_text(prompt):
 
 
 
14
  try:
15
- inputs = tokenizer(prompt, return_tensors="pt")
16
- outputs = model.generate(**inputs, max_length=100)
17
- return tokenizer.decode(outputs[0], skip_special_tokens=True)
18
  except Exception as e:
19
  return f"Ошибка при генерации текста: {str(e)}"
20
 
 
1
+ from transformers import pipeline
 
 
2
 
3
  # Имя модели
4
  model_name = "nvidia/Hymba-1.5B-Instruct"
5
 
6
+ # Создание пайплайна
7
+ pipe = pipeline("text-generation", model=model_name, trust_remote_code=True)
 
8
 
9
  # Функция для генерации текста
10
  def generate_text(prompt):
11
+ messages = [
12
+ {"role": "user", "content": prompt},
13
+ ]
14
  try:
15
+ response = pipe(messages)[0]["generated_text"]
16
+ return response
 
17
  except Exception as e:
18
  return f"Ошибка при генерации текста: {str(e)}"
19