lukmanaj commited on
Commit
46c1cbd
·
verified ·
1 Parent(s): c2f387f

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +47 -16
app.py CHANGED
@@ -7,8 +7,6 @@ import pandas as pd
7
  # from google.genai import types
8
  import torch
9
  from transformers import AutoModelForCausalLM, AutoTokenizer
10
- from smolagents.agents import ReActAgent
11
- from smolagents.tools import tool
12
 
13
  # (Keep Constants as is)
14
  # --- Constants ---
@@ -69,39 +67,72 @@ DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
69
  # print(f"Error during Gemini API call: {str(e)}")
70
  # return f"Error: {str(e)}"
71
 
72
- class BasicAgent(ReActAgent):
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
73
  def __init__(self):
74
  print("BasicAgent using local LLM initialized.")
75
 
76
- # Load a small model from Hugging Face
77
- model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0" # You can pick another lightweight model
78
  self.tokenizer = AutoTokenizer.from_pretrained(model_name)
79
  self.model = AutoModelForCausalLM.from_pretrained(
80
  model_name,
81
  torch_dtype=torch.float16,
82
- device_map="auto" # Automatically choose GPU/CPU
83
  )
84
 
85
- super().__init__(tools=[]) # No tools for now
86
-
87
- def call(self, task: str) -> str:
88
- """Core method for answering a task."""
89
- prompt = f"Answer the following question concisely:\n\n{task}\n\nAnswer:"
90
  inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
91
 
92
  with torch.no_grad():
93
  outputs = self.model.generate(
94
  **inputs,
95
- max_new_tokens=200,
96
  do_sample=True,
97
  temperature=0.7,
98
- top_p=0.95,
99
  top_k=50,
100
  )
101
- answer = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
102
 
103
- # Extract only the answer part
104
- return answer.split("Answer:")[-1].strip()
 
 
105
 
106
  def run_and_submit_all( profile: gr.OAuthProfile | None):
107
  """
 
7
  # from google.genai import types
8
  import torch
9
  from transformers import AutoModelForCausalLM, AutoTokenizer
 
 
10
 
11
  # (Keep Constants as is)
12
  # --- Constants ---
 
67
  # print(f"Error during Gemini API call: {str(e)}")
68
  # return f"Error: {str(e)}"
69
 
70
+ # class BasicAgent(ReActAgent):
71
+ # def __init__(self):
72
+ # print("BasicAgent using local LLM initialized.")
73
+
74
+ # # Load a small model from Hugging Face
75
+ # model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0" # You can pick another lightweight model
76
+ # self.tokenizer = AutoTokenizer.from_pretrained(model_name)
77
+ # self.model = AutoModelForCausalLM.from_pretrained(
78
+ # model_name,
79
+ # torch_dtype=torch.float16,
80
+ # device_map="auto" # Automatically choose GPU/CPU
81
+ # )
82
+
83
+ # super().__init__(tools=[]) # No tools for now
84
+
85
+ # def call(self, task: str) -> str:
86
+ # """Core method for answering a task."""
87
+ # prompt = f"Answer the following question concisely:\n\n{task}\n\nAnswer:"
88
+ # inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
89
+
90
+ # with torch.no_grad():
91
+ # outputs = self.model.generate(
92
+ # **inputs,
93
+ # max_new_tokens=200,
94
+ # do_sample=True,
95
+ # temperature=0.7,
96
+ # top_p=0.95,
97
+ # top_k=50,
98
+ # )
99
+ # answer = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
100
+
101
+ # # Extract only the answer part
102
+ # return answer.split("Answer:")[-1].strip()
103
+ class BasicAgent:
104
  def __init__(self):
105
  print("BasicAgent using local LLM initialized.")
106
 
107
+ # Load a small Hugging Face model
108
+ model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0" # Change if you want
109
  self.tokenizer = AutoTokenizer.from_pretrained(model_name)
110
  self.model = AutoModelForCausalLM.from_pretrained(
111
  model_name,
112
  torch_dtype=torch.float16,
113
+ device_map="auto" # Use GPU if available
114
  )
115
 
116
+ def __call__(self, task: str) -> str:
117
+ """Answer a question."""
118
+ prompt = f"Answer the following question clearly and concisely:\n\n{task}\n\nAnswer:"
 
 
119
  inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
120
 
121
  with torch.no_grad():
122
  outputs = self.model.generate(
123
  **inputs,
124
+ max_new_tokens=256,
125
  do_sample=True,
126
  temperature=0.7,
127
+ top_p=0.9,
128
  top_k=50,
129
  )
130
+ decoded = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
131
 
132
+ # Extract the answer part
133
+ if "Answer:" in decoded:
134
+ return decoded.split("Answer:")[-1].strip()
135
+ return decoded.strip()
136
 
137
  def run_and_submit_all( profile: gr.OAuthProfile | None):
138
  """