arya-narrable commited on
Commit
ed4f7f8
·
1 Parent(s): ce2562b
Files changed (1) hide show
  1. app.py +22 -9
app.py CHANGED
@@ -1,15 +1,28 @@
1
  import gradio as gr
2
  from transformers import AutoModelForCausalLM, AutoTokenizer
3
 
4
- model_name = "internistai/base-7b-v0.2"
5
- tokenizer = AutoTokenizer.from_pretrained(model_name)
6
- model = AutoModelForCausalLM.from_pretrained(model_name)
 
 
 
 
 
 
7
 
8
- def chat(message):
9
- inputs = tokenizer(message, return_tensors="pt")
10
- outputs = model.generate(**inputs)
11
- response = tokenizer.decode(outputs[0], skip_special_tokens=True)
12
- return response
 
 
 
13
 
14
- iface = gr.Interface(fn=chat, inputs="text", outputs="text")
 
 
 
 
15
  iface.launch()
 
1
  import gradio as gr
2
  from transformers import AutoModelForCausalLM, AutoTokenizer
3
 
4
+ def load_model():
5
+ try:
6
+ model_name = "internistai/base-7b-v0.2"
7
+ tokenizer = AutoTokenizer.from_pretrained(model_name)
8
+ model = AutoModelForCausalLM.from_pretrained(model_name)
9
+ return model, tokenizer
10
+ except Exception as e:
11
+ print(f"Error loading model: {e}")
12
+ return None, None
13
 
14
+ def chat(message, model, tokenizer):
15
+ try:
16
+ inputs = tokenizer(message, return_tensors="pt")
17
+ outputs = model.generate(**inputs)
18
+ response = tokenizer.decode(outputs[0], skip_special_tokens=True)
19
+ return response
20
+ except Exception as e:
21
+ return f"Error generating response: {e}"
22
 
23
+ model, tokenizer = load_model()
24
+ if model is None or tokenizer is None:
25
+ print("Failed to load model or tokenizer. Please check the configuration.")
26
+
27
+ iface = gr.Interface(fn=lambda message: chat(message, model, tokenizer), inputs="text", outputs="text")
28
  iface.launch()